codingutf8 2018 antonio mendoza hao tan mohit bansal adapted from facebook inc detectron2 huggingface co 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 import copy other location will always be cpu for np tensors convert to boxes to use the clip function select max scores apply nms helper functions args proposals listtensor l n hiwia 4 predobjectnesslogits tensors of length l nmsthresh float iou threshold to use for nms prenmstopk int before nms postnmstopk int after nms minboxsidelen float minimum proposal box side training bool true if proposals are to be used in training returns results listdict stores postnmstopk object proposals for image i 1 select topk anchor for every level and every image sort is faster than topk https github compytorchpytorchissues22812 topkscoresi topkidx logitsi topknumproposalsi dim1 each is n x topk 2 concat all levels together if i change to batchednms i wonder if this will make a difference 3 for each image run a perlevel nms and choose topk results i will have to take a look at the boxes clip method filter empty boxes i wonder if it would be possible for me to pad all these things returns posidx negidx tensor 1d vector of indices the total length of both is numsamples or fewer protect against not enough positive examples protect against not enough negative examples randomly select positive and negative examples eqn 1 in fpn paper clamp level to min max in case the box size is too large or too small for the available feature maps helper classes this rcnn transformation scales the box s width and height by expdw expdh and shifts a box s center by the offset dx width dy height args weights 4element tuple scaling factors that are applied to the dx dy dw dh deltas in fast rcnn these were originally set such that the deltas have unit variance now they are treated as hyperparameters of the system scaleclamp float when predicting deltas the predicted box scaling factors dw and dh are clamped such that they are scaleclamp value for clamping large dw and dh predictions the heuristic is that we clamp such that dw and dh are no larger than what would transform a 16px box into a 1000px box based on a small anchor 16px and a typical image size 1000px get box regression transformation deltas dx dy dw dh that can be used to transform the srcboxes into the targetboxes that is the relation targetboxes self applydeltasdeltas srcboxes is true unless any delta is too large and is clamped args srcboxes tensor source boxes e g object proposals targetboxes tensor target of the transformation e g groundtruth boxes apply transformation deltas dx dy dw dh to boxes args deltas tensor transformation deltas of shape n k4 where k 1 deltasi represents k potentially different classspecific box transformations for the single box boxesi boxes tensor boxes to transform of shape n 4 prevent sending too large values into torch exp this class assigns to each predicted element e g a box a groundtruth element each predicted element will have exactly zero or one matches each groundtruth element may be matched to zero or more predicted elements the matching is determined by the mxn matchqualitymatrix that characterizes how well each groundtruth predictionpair match each other for example if the elements are boxes this matrix may contain box intersectionoverunion overlap values the matcher returns a a vector of length n containing the index of the groundtruth element m in 0 m that matches to prediction n in 0 n b a vector of length n containing the labels for each prediction args thresholds list a list of thresholds used to stratify predictions into levels labels list a list of values to label predictions belonging at each level a label can be one of 1 0 1 signifying ignore negative class positive class respectively allowlowqualitymatches bool if true produce additional matches or predictions with maximum match quality lower than highthreshold for example thresholds 0 3 0 5 labels 0 1 1 all predictions with iou 0 3 will be marked with 0 and thus will be considered as false positives while training all predictions with 0 3 iou 0 5 will be marked with 1 and thus will be ignored all predictions with 0 5 iou will be marked with 1 and thus will be considered as true positives args matchqualitymatrix tensorfloat an mxn tensor containing the pairwise quality between m groundtruth elements and n predicted elements all elements must be 0 due to the us of torch nonzero for selecting indices in meth setlowqualitymatches returns matches tensorint64 a vector of length n where matchesi is a matched groundtruth index in 0 m matchlabels tensorint8 a vector of length n where predlabelsi indicates true or false positive or ignored when no gt boxes exist we define iou 0 and therefore set labels to self labels0 which usually defaults to background class 0 to choose to ignore instead can make labels1 0 1 1 set appropriate thresholds matchqualitymatrix is m gt x n predicted max over gt elements dim 0 to find best gt candidate for each prediction produce additional matches for predictions that have only lowquality matches specifically for each groundtruth g find the set of predictions that have maximum overlap with it including ties for each prediction in that set if it is unmatched then match it to the groundtruth g this function implements the rpn assignment case i in sec 3 1 2 of faster rcnn for each gt find the prediction with which it has highest quality find the highest quality match available even if it is low including ties note that the matches qualities must be positive due to the use of torch nonzero args box2boxtransform box2boxtransform class box2boxtransform instance for anchorproposal transformations anchormatcher matcher class matcher instance for matching anchors to groundtruth boxes used to determine training labels batchsizeperimage int number of proposals to sample when training positivefraction float target fraction of sampled proposals that should be positive images imagelist class imagelist instance representing n input images predobjectnesslogits listtensor a list of l elements element i is a tensor of shape n a hi w predanchordeltas listtensor a list of l elements element i is a tensor of shape n a4 hi wi anchors listtorch tensor nested list of boxes anchorsij at n l stores anchor array for feature map l boundarythreshold int if 0 then anchors that extend beyond the image boundary by more than boundarythresh are not used in training gtboxes listboxes optional a list of n elements smoothl1beta float the transition point between l1 and l2 lossn when set to 0 the loss becomes l1 when inf it is ignored predanchordeltas l n hi wi anchors n l 1 b here we loop over specific feature map not images append feature map proposals with shape n hiwia b returns predobjectnesslogits listtensor n hiwia reshape n a hi wi n hi wi a n hiwia main classes this module is used in the original fpn to generate a downsampled p6 feature from p5 this module is used in retinanet to generate extra layers p6 and p7 from c5 feature use pad 1 instead of pad zero the original msra resnet models have stride in the first 1x1 conv the subsequent fb torch resnet and caffe2 resnext implementations have stride in the 3x3 conv some backbones require the input height and width to be divisible by a specific integer this is typically true for encoder decoder type networks with lateral connection e g fpn for which feature maps need to match dimension in the bottom up and top down paths set to 0 if no specific input size divisibility is required deprecated return self outfeatures property def outfeaturestridesself deprecated return f self outfeaturechannelsf for f in self outfeatures class resnetbackbone def initself stem stages numclassesnone outfeaturesnone superresnet self init self stem stem self numclasses numclasses currentstride self stem stride self outfeaturestrides stem currentstride self outfeaturechannels stem self stem outchannels self stagesandnames for i blocks in enumeratestages for block in blocks assert isinstanceblock resnetblockbase block currchannels block outchannels stage nn sequentialblocks name res stri 2 self addmodulename stage self stagesandnames appendstage name self outfeaturestridesname currentstride int currentstride np prodk stride for k in blocks self outfeaturechannelsname blocks1 outchannels if numclasses is not none self avgpool nn adaptiveavgpool2d1 1 self linear nn linearcurrchannels numclasses sec 5 1 in accurate large minibatch sgd training imagenet in 1 hour the 1000way fullyconnected layer is initialized by drawing weights from a zeromean gaussian with std of 0 01 nn init normalself linear weight stddev0 01 name linear if outfeatures is none outfeatures name self outfeatures outfeatures assert lenself outfeatures children x0 for x in self namedchildren for outfeature in self outfeatures assert outfeature in children available children format joinchildren def forwardself x outputs x self stemx if stem in self outfeatures outputsstem x for stage name in self stagesandnames x stagex if name in self outfeatures outputsname x if self numclasses is not none x self avgpoolx x self linearx if linear in self outfeatures outputslinear x return outputs def outputshapeself return name shapespec channelsself outfeaturechannelsname strideself outfeaturestridesname for name in self outfeatures staticmethod def makestage blockclass numblocks firststridenone inchannels outchannels kwargs if firststride is not none assert stride not in kwargs and strideperblock not in kwargs kwargsstrideperblock firststride 1 numblocks 1 blocks for i in rangenumblocks currkwargs for k v in kwargs items if k endswithperblock assert lenv numblocks fargument k of makestage should have the same length as numblocksnumblocks newk k lenperblock assert newk not in kwargs fcannot call makestage with both k and newk currkwargsnewk vi else currkwargsk v blocks appendblockclassinchannelsinchannels outchannelsoutchannels currkwargs inchannels outchannels return blocks class roipoolernn module def init self outputsize scales samplingratio canonicalboxsize224 canonicallevel4 super init assumption that stride is a power of 2 minlevel math log2scales0 maxlevel math log2scales1 a bunch of testing assert math iscloseminlevel intminlevel and math isclosemaxlevel intmaxlevel assert lenscales maxlevel minlevel 1 not pyramid assert 0 minlevel and minlevel maxlevel if isinstanceoutputsize int outputsize outputsize outputsize assert lenoutputsize 2 and isinstanceoutputsize0 int and isinstanceoutputsize1 int if lenscales 1 assert minlevel canonicallevel and canonicallevel maxlevel assert canonicalboxsize 0 self outputsize outputsize self minlevel intminlevel self maxlevel intmaxlevel self levelpoolers nn modulelistroipooloutputsize spatialscalescale for scale in scales self canonicallevel canonicallevel self canonicalboxsize canonicalboxsize def forwardself featuremaps boxes x listfeaturemaps values numlevelassignments lenself levelpoolers assert lenx numlevelassignments and lenboxes x0 size0 poolerfmtboxes convertboxestopoolerformatboxes if numlevelassignments 1 return self levelpoolers0x0 poolerfmtboxes levelassignments assignboxestolevels boxes self minlevel self maxlevel self canonicalboxsize self canonicallevel numboxes lenpoolerfmtboxes numchannels x0 shape1 outputsize self outputsize0 dtype device x0 dtype x0 device output torch zeros numboxes numchannels outputsize outputsize dtypedtype devicedevice for level xlevel pooler in enumeratezipx self levelpoolers inds torch nonzerolevelassignments level squeeze1 poolerfmtboxeslevel poolerfmtboxesinds outputinds poolerxlevel poolerfmtboxeslevel return output class roioutputsobject def initself cfg trainingfalse self smoothl1beta cfg roiboxhead smoothl1beta self box2boxtransform box2boxtransformweightscfg roiboxhead bboxregweights self training training self scorethresh cfg roiheads scorethreshtest self mindetections cfg mindetections self maxdetections cfg maxdetections nmsthresh cfg roiheads nmsthreshtest if not isinstancenmsthresh list nmsthresh nmsthresh self nmsthresh nmsthresh def predictboxesself proposals boxdeltas predsperimage numpred boxdeltas size0 b proposals0 size1 k boxdeltas size1 b boxdeltas boxdeltas viewnumpred k b proposals torch catproposals dim0 unsqueeze2 expandnumpred k b proposals proposals reshape1 b boxes self box2boxtransform applydeltasboxdeltas proposals return boxes viewnumpred k b splitpredsperimage dim0 def predictobjsself objlogits predsperimage probs nn functional softmaxobjlogits dim1 probs probs splitpredsperimage dim0 return probs def predictattrsself attrlogits predsperimage attrlogits attrlogits 1 softmax1 attrprobs attrs attrlogits max1 return attrprobs splitpredsperimage dim0 attrs splitpredsperimage dim0 torch nograd def inference self objlogits attrlogits boxdeltas predboxes features sizes scalesnone only the pred boxes is the predsperimage p size0 for p in predboxes boxesall self predictboxespredboxes boxdeltas predsperimage objscoresall self predictobjsobjlogits predsperimage list of length n attrprobsall attrsall self predictattrsattrlogits predsperimage features features splitpredsperimage dim0 fun for each image too also i can experiment and do multiple images finalresults zipped zipboxesall objscoresall attrprobsall attrsall sizes for i boxes objscores attrprobs attrs size in enumeratezipped for nmst in self nmsthresh outputs donms boxes objscores size self scorethresh nmst self mindetections self maxdetections if outputs is not none maxboxes maxscores classes ids outputs break if scales is not none scaleyx scalesi maxboxes 0 2 scaleyx1 maxboxes 1 2 scaleyx0 finalresults append maxboxes classes maxscores attrsids attrprobsids featuresiids boxes classes classprobs attrs attrprobs roifeatures maplist zipfinalresults return boxes classes classprobs attrs attrprobs roifeatures def trainingself objlogits attrlogits boxdeltas predboxes features sizes pass def call self objlogits attrlogits boxdeltas predboxes features sizes scalesnone if self training raise notimplementederror return self inference objlogits attrlogits boxdeltas predboxes features sizes scalesscales class res5roiheadsnn module def initself cfg inputshape super init self batchsizeperimage cfg rpn batchsizeperimage self positivesamplefraction cfg roiheads positivefraction self infeatures cfg roiheads infeatures self numclasses cfg roiheads numclasses self proposalappendgt cfg roiheads proposalappendgt self featurestrides k v stride for k v in inputshape items self featurechannels k v channels for k v in inputshape items self clsagnosticbboxreg cfg roiboxhead clsagnosticbboxreg self stagechannelfactor 23 res5 is 8x res2 self outchannels cfg resnets res2outchannels self stagechannelfactor self proposalmatcher matcher cfg roiheads iouthresholds cfg roiheads ioulabels allowlowqualitymatchesfalse poolerresolution cfg roiboxhead poolerresolution poolerscales 1 0 self featurestridesself infeatures0 samplingratio cfg roiboxhead poolersamplingratio res5halve cfg roiboxhead res5halve useattr cfg roiboxhead attr numattrs cfg roiboxhead numattrs self pooler roipooler outputsizepoolerresolution scalespoolerscales samplingratiosamplingratio self res5 self buildres5blockcfg if not res5halve self res50 conv1 stride 1 1 self res50 shortcut stride 1 1 for i in range3 self res5i conv2 padding 2 2 self res5i conv2 dilation 2 2 self boxpredictor fastrcnnoutputlayers self outchannels self numclasses self clsagnosticbboxreg useattruseattr numattrsnumattrs def buildres5blockself cfg stagechannelfactor self stagechannelfactor res5 is 8x res2 numgroups cfg resnets numgroups widthpergroup cfg resnets widthpergroup bottleneckchannels numgroups widthpergroup stagechannelfactor outchannels self outchannels stridein1x1 cfg resnets stridein1x1 norm cfg resnets norm blocks resnet makestage bottleneckblock 3 firststride2 inchannelsoutchannels 2 bottleneckchannelsbottleneckchannels outchannelsoutchannels numgroupsnumgroups normnorm stridein1x1stridein1x1 return nn sequentialblocks def sharedroitransformself features boxes x self poolerfeatures boxes return self res5x def forwardself features proposalboxes gtboxesnone if self training raise notimplementederror assert not proposalboxes0 requiresgrad boxfeatures self sharedroitransformfeatures proposalboxes featurepooled boxfeatures meandim2 3 pooled to 1x1 objlogits attrlogits predproposaldeltas self boxpredictorfeaturepooled return objlogits attrlogits predproposaldeltas featurepooled class anchorgeneratornn module def initself cfg inputshape listshapespec super init sizes cfg anchorgenerator sizes aspectratios cfg anchorgenerator aspectratios self strides x stride for x in inputshape self offset cfg anchorgenerator offset assert 0 0 self offset 1 0 self offset self numfeatures lenself strides self cellanchors nn parameterlistself calculateanchorssizes aspectratios self spacialfeatdim 4 def calculateanchorsself sizes aspectratios if one size or aspect ratio is specified and there are multiple feature maps then we broadcast anchors of that single size or aspect ratio if lensizes 1 sizes self numfeatures if lenaspectratios 1 aspectratios self numfeatures assert self numfeatures lensizes assert self numfeatures lenaspectratios cellanchors self generatecellanchorss a float for s a in zipsizes aspectratios return cellanchors property def boxdimself return self spacialfeatdim property def numcellanchorsself return lencellanchors for cellanchors in self cellanchors def gridanchorsself gridsizes anchors for size stride baseanchors in zipgridsizes self strides self cellanchors shiftx shifty creategridoffsetssize stride self offset baseanchors device shifts torch stackshiftx shifty shiftx shifty dim1 anchors appendshifts view1 1 4 baseanchors view1 1 4 reshape1 4 return anchors def generatecellanchorsself sizes32 64 128 256 512 aspectratios0 5 1 2 anchors for size in sizes area size2 0 for aspectratio in aspectratios w math sqrtarea aspectratio h aspectratio w x0 y0 x1 y1 w 2 0 h 2 0 w 2 0 h 2 0 anchors appendx0 y0 x1 y1 return nn parametertorch tensoranchors def forwardself features numimages features0 size0 gridsizes featuremap shape2 for featuremap in features anchorsoverallfeaturemaps self gridanchorsgridsizes anchorsoverallfeaturemaps torch stackanchorsoverallfeaturemaps return anchorsoverallfeaturemaps unsqueeze0 repeatinterleavenumimages dim0 class rpnheadnn module def initself cfg inputshape listshapespec super init standard rpn is shared across levels inchannels s channels for s in inputshape assert lensetinchannels 1 each level must have the same channel inchannels inchannels0 anchorgenerator anchorgeneratorcfg inputshape numcellanchors anchorgenerator numcellanchors boxdim anchorgenerator boxdim assert lensetnumcellanchors 1 each level must have the same number of cell anchors numcellanchors numcellanchors0 if cfg proposalgenerator hiddenchannels 1 hidchannels inchannels else hidchannels cfg proposalgenerator hiddenchannels modifications for vg in rpn modelingproposalgeneratorrpn py use hidden dim instead fo the same dim as res4 inchannels 3x3 conv for the hidden representation self conv nn conv2dinchannels hidchannels kernelsize3 stride1 padding1 1x1 conv for predicting objectness logits self objectnesslogits nn conv2dhidchannels numcellanchors kernelsize1 stride1 1x1 conv for predicting box2box transform deltas self anchordeltas nn conv2dhidchannels numcellanchors boxdim kernelsize1 stride1 for layer in self conv self objectnesslogits self anchordeltas nn init normallayer weight std0 01 nn init constantlayer bias 0 def forwardself features predobjectnesslogits predanchordeltas for x in features t nn functional reluself convx predobjectnesslogits appendself objectnesslogitst predanchordeltas appendself anchordeltast return predobjectnesslogits predanchordeltas class rpnnn module def initself cfg inputshape dictstr shapespec super init self minboxsidelen cfg proposalgenerator minsize self infeatures cfg rpn infeatures self nmsthresh cfg rpn nmsthresh self batchsizeperimage cfg rpn batchsizeperimage self positivefraction cfg rpn positivefraction self smoothl1beta cfg rpn smoothl1beta self lossweight cfg rpn lossweight self prenmstopk true cfg rpn prenmstopktrain false cfg rpn prenmstopktest self postnmstopk true cfg rpn postnmstopktrain false cfg rpn postnmstopktest self boundarythreshold cfg rpn boundarythresh self anchorgenerator anchorgeneratorcfg inputshapef for f in self infeatures self box2boxtransform box2boxtransformweightscfg rpn bboxregweights self anchormatcher matcher cfg rpn iouthresholds cfg rpn ioulabels allowlowqualitymatchestrue self rpnhead rpnheadcfg inputshapef for f in self infeatures def trainingself images imageshapes features gtboxes pass def inferenceself outputs images imageshapes features gtboxesnone outputs findtoprpnproposals outputs predictproposals outputs predictobjectnesslogits images imageshapes self nmsthresh self prenmstopkself training self postnmstopkself training self minboxsidelen self training results for img in outputs imboxes imgboxlogits img imgboxlogits inds imgboxlogits sortdescendingtrue imboxes imboxesinds results appendimboxes imgboxlogits proposalboxes logits tuplemaplist zipresults return proposalboxes logits def forwardself images imageshapes features gtboxesnone features is dict key block level v featuremap features featuresf for f in self infeatures predobjectnesslogits predanchordeltas self rpnheadfeatures anchors self anchorgeneratorfeatures outputs rpnoutputs self box2boxtransform self anchormatcher self batchsizeperimage self positivefraction images predobjectnesslogits predanchordeltas anchors self boundarythreshold gtboxes self smoothl1beta for rpnonly models the proposals are the final output if self training raise notimplementederror return self trainingoutputs images imageshapes features gtboxes else return self inferenceoutputs images imageshapes features gtboxes class fastrcnnoutputlayersnn module def init self inputsize numclasses clsagnosticbboxreg boxdim4 useattrfalse numattrs1 super init if not isinstanceinputsize int inputsize np prodinputsize do 1 for background class self clsscore nn linearinputsize numclasses 1 numbboxregclasses 1 if clsagnosticbboxreg else numclasses self bboxpred nn linearinputsize numbboxregclasses boxdim self useattr useattr if useattr self clsembedding nn embeddingnumclasses 1 inputsize 8 self fcattr nn linearinputsize inputsize 8 inputsize 4 self attrscore nn linearinputsize 4 numattrs 1 nn init normalself clsscore weight std0 01 nn init normalself bboxpred weight std0 001 for item in self clsscore self bboxpred nn init constantitem bias 0 def forwardself roifeatures if roifeatures dim 2 roifeatures torch flattenroifeatures startdim1 scores self clsscoreroifeatures proposaldeltas self bboxpredroifeatures if self useattr maxclass scores max1 b c b clsemb self clsembeddingmaxclass b b 256 roifeatures torch catroifeatures clsemb 1 b 2048 b 256 b 2304 roifeatures self fcattrroifeatures roifeatures nn functional reluroifeatures attrscores self attrscoreroifeatures return scores attrscores proposaldeltas else return scores proposaldeltas class generalizedrcnnnn module def initself cfg super init self device torch devicecfg model device self backbone buildbackbonecfg self proposalgenerator rpncfg self backbone outputshape self roiheads res5roiheadscfg self backbone outputshape self roioutputs roioutputscfg self toself device classmethod def frompretrainedcls pretrainedmodelnameorpath modelargs kwargs config kwargs popconfig none statedict kwargs popstatedict none cachedir kwargs popcachedir none fromtf kwargs popfromtf false forcedownload kwargs popforcedownload false resumedownload kwargs popresumedownload false proxies kwargs popproxies none localfilesonly kwargs poplocalfilesonly false usecdn kwargs popusecdn true load config if we don t provide a configuration if not isinstanceconfig config configpath config if config is not none else pretrainedmodelnameorpath try config config frompretrained configpath cachedircachedir forcedownloadforcedownload resumedownloadresumedownload proxiesproxies localfilesonlylocalfilesonly load model if pretrainedmodelnameorpath is not none if os path isdirpretrainedmodelnameorpath if os path isfileos path joinpretrainedmodelnameorpath weightsname load from a pytorch checkpoint archivefile os path joinpretrainedmodelnameorpath weightsname else raise environmenterror error no file named found in directory format weightsname pretrainedmodelnameorpath elif os path isfilepretrainedmodelnameorpath or isremoteurlpretrainedmodelnameorpath archivefile pretrainedmodelnameorpath elif os path isfilepretrainedmodelnameorpath index assert fromtf we found a tensorflow checkpoint at please set fromtf to true to load from this checkpoint format pretrainedmodelnameorpath index archivefile pretrainedmodelnameorpath index else archivefile hfbucketurl pretrainedmodelnameorpath filenameweightsname usecdnusecdn try load from url or cache if already cached resolvedarchivefile cachedpath archivefile cachedircachedir forcedownloadforcedownload proxiesproxies resumedownloadresumedownload localfilesonlylocalfilesonly if resolvedarchivefile is none raise environmenterror except environmenterror msg fcan t load weights for pretrainedmodelnameorpath raise environmenterrormsg if resolvedarchivefile archivefile printloading weights file formatarchivefile else printloading weights file from cache at formatarchivefile resolvedarchivefile else resolvedarchivefile none instantiate model model clsconfig if statedict is none try try statedict torch loadresolvedarchivefile maplocationcpu except exception statedict loadcheckpointresolvedarchivefile except exception raise oserror unable to load weights from pytorch checkpoint file if you tried to load a pytorch model from a tf 2 0 checkpoint please set fromtftrue missingkeys unexpectedkeys errormsgs convert old format to new format if needed from a pytorch statedict oldkeys newkeys for key in statedict keys newkey none if gamma in key newkey key replacegamma weight if beta in key newkey key replacebeta bias if newkey oldkeys appendkey newkeys appendnewkey for oldkey newkey in zipoldkeys newkeys statedictnewkey statedict popoldkey copy statedict so loadfromstatedict can modify it metadata getattrstatedict metadata none statedict statedict copy if metadata is not none statedict metadata metadata modeltoload model modeltoload loadstatedictstatedict if model class name modeltoload class name basemodelstatedict modeltoload statedict keys headmodelstatedictwithoutbaseprefix key splitcls basemodelprefix 1 for key in model statedict keys missingkeys extendheadmodelstatedictwithoutbaseprefix basemodelstatedict if lenunexpectedkeys 0 print fsome weights of the model checkpoint at pretrainedmodelnameorpath were not used when f initializing model class name unexpectedkeysn this is expected if you are f initializing model class name from the checkpoint of a model trained on another task or with another architecture e g initializing a bertforsequenceclassification model from a bertforpretraining model n this is not expected if you are initializing f model class name from the checkpoint of a model that you expect to be exactly identical initializing a bertforsequenceclassification model from a bertforsequenceclassification model else printfall model checkpoint weights were used when initializing model class name n if lenmissingkeys 0 print fsome weights of model class name were not initialized from the model checkpoint at f pretrainedmodelnameorpath and are newly initialized missingkeysnyou should probably train this model on a downstream task to be able to use it for predictions and inference else print fall the weights of model class name were initialized from the model checkpoint at f pretrainedmodelnameorpath nif your task is similar to the task the model of the checkpoint f was trained on you can already use model class name for predictions without further training if lenerrormsgs 0 raise runtimeerror errors in loading statedict for nt format model class name nt joinerrormsgs set model in evaluation mode to deactivate dropout modules by default model eval return model def forward self images imageshapes gtboxesnone proposalsnone scalesyxnone kwargs if self training raise notimplementederror return self inference imagesimages imageshapesimageshapes gtboxesgtboxes proposalsproposals scalesyxscalesyx kwargs torch nograd def inference self images imageshapes gtboxesnone proposalsnone scalesyxnone kwargs run images through backbone originalsizes imageshapes scalesyx features self backboneimages generate proposals if none are available if proposals is none proposalboxes self proposalgeneratorimages imageshapes features gtboxes else assert proposals is not none pool object features from either gtboxes or from proposals objlogits attrlogits boxdeltas featurepooled self roiheadsfeatures proposalboxes gtboxes prepare frcnn outputs and select top proposals boxes classes classprobs attrs attrprobs roifeatures self roioutputs objlogitsobjlogits attrlogitsattrlogits boxdeltasboxdeltas predboxesproposalboxes featuresfeaturepooled sizesimageshapes scalesscalesyx will we pad subsetkwargs maxdetections kwargs getmaxdetections none returntensors kwargs getreturntensors none padvalue kwargs getpadvalue 0 padding kwargs getpadding none predsperimage torch tensorp size0 for p in boxes boxes padlisttensorsboxes predsperimage subsetkwargs classes padlisttensorsclasses predsperimage subsetkwargs classprobs padlisttensorsclassprobs predsperimage subsetkwargs attrs padlisttensorsattrs predsperimage subsetkwargs attrprobs padlisttensorsattrprobs predsperimage subsetkwargs roifeatures padlisttensorsroifeatures predsperimage subsetkwargs subsetkwargspadding none predsperimage padlisttensorspredsperimage none subsetkwargs sizes padlisttensorsimageshapes none subsetkwargs normalizedboxes normboxboxes originalsizes return ordereddict objids classes objprobs classprobs attrids attrs attrprobs attrprobs boxes boxes sizes sizes predsperimage predsperimage roifeatures roifeatures normalizedboxes normalizedboxes other location will always be cpu for np tensors convert to boxes to use the clip function r x c x 4 select max scores r x c r select max boxes according to the max scores apply nms helper functions keep for debugging args proposals list tensor l n hi wi a 4 pred_objectness_logits tensors of length l nms_thresh float iou threshold to use for nms pre_nms_topk int before nms post_nms_topk int after nms min_box_side_len float minimum proposal box side training bool true if proposals are to be used in training returns results list dict stores post_nms_topk object proposals for image i 1 select top k anchor for every level and every image lvl tensor each of shape n x topk lvl tensor each of shape topk sort is faster than topk https github com pytorch pytorch issues 22812 topk_scores_i topk_idx logits_i topk num_proposals_i dim 1 each is n x topk n x topk x 4 2 concat all levels together if i change to batched_nms i wonder if this will make a difference 3 for each image run a per level nms and choose topk results i will have to take a look at the boxes clip method filter empty boxes i wonder if it would be possible for me to pad all these things returns pos_idx neg_idx tensor 1d vector of indices the total length of both is num_samples or fewer protect against not enough positive examples protect against not enough negative examples randomly select positive and negative examples eqn 1 in fpn paper clamp level to min max in case the box size is too large or too small for the available feature maps helper classes this r cnn transformation scales the box s width and height by exp dw exp dh and shifts a box s center by the offset dx width dy height args weights 4 element tuple scaling factors that are applied to the dx dy dw dh deltas in fast r cnn these were originally set such that the deltas have unit variance now they are treated as hyperparameters of the system scale_clamp float when predicting deltas the predicted box scaling factors dw and dh are clamped such that they are scale_clamp value for clamping large dw and dh predictions the heuristic is that we clamp such that dw and dh are no larger than what would transform a 16px box into a 1000px box based on a small anchor 16px and a typical image size 1000px get box regression transformation deltas dx dy dw dh that can be used to transform the src_boxes into the target_boxes that is the relation target_boxes self apply_deltas deltas src_boxes is true unless any delta is too large and is clamped args src_boxes tensor source boxes e g object proposals target_boxes tensor target of the transformation e g ground truth boxes apply transformation deltas dx dy dw dh to boxes args deltas tensor transformation deltas of shape n k 4 where k 1 deltas i represents k potentially different class specific box transformations for the single box boxes i boxes tensor boxes to transform of shape n 4 prevent sending too large values into torch exp x1 y1 x2 y2 this class assigns to each predicted element e g a box a ground truth element each predicted element will have exactly zero or one matches each ground truth element may be matched to zero or more predicted elements the matching is determined by the mxn match_quality_matrix that characterizes how well each ground truth prediction pair match each other for example if the elements are boxes this matrix may contain box intersection over union overlap values the matcher returns a a vector of length n containing the index of the ground truth element m in 0 m that matches to prediction n in 0 n b a vector of length n containing the labels for each prediction args thresholds list a list of thresholds used to stratify predictions into levels labels list a list of values to label predictions belonging at each level a label can be one of 1 0 1 signifying ignore negative class positive class respectively allow_low_quality_matches bool if true produce additional matches or predictions with maximum match quality lower than high_threshold for example thresholds 0 3 0 5 labels 0 1 1 all predictions with iou 0 3 will be marked with 0 and thus will be considered as false positives while training all predictions with 0 3 iou 0 5 will be marked with 1 and thus will be ignored all predictions with 0 5 iou will be marked with 1 and thus will be considered as true positives args match_quality_matrix tensor float an mxn tensor containing the pairwise quality between m ground truth elements and n predicted elements all elements must be 0 due to the us of torch nonzero for selecting indices in meth set_low_quality_matches_ returns matches tensor int64 a vector of length n where matches i is a matched ground truth index in 0 m match_labels tensor int8 a vector of length n where pred_labels i indicates true or false positive or ignored when no gt boxes exist we define iou 0 and therefore set labels to self labels 0 which usually defaults to background class 0 to choose to ignore instead can make labels 1 0 1 1 set appropriate thresholds match_quality_matrix is m gt x n predicted max over gt elements dim 0 to find best gt candidate for each prediction produce additional matches for predictions that have only low quality matches specifically for each ground truth g find the set of predictions that have maximum overlap with it including ties for each prediction in that set if it is unmatched then match it to the ground truth g this function implements the rpn assignment case i in sec 3 1 2 of faster r cnn for each gt find the prediction with which it has highest quality find the highest quality match available even if it is low including ties note that the matches qualities must be positive due to the use of torch nonzero args box2box_transform box2boxtransform class box2boxtransform instance for anchor proposal transformations anchor_matcher matcher class matcher instance for matching anchors to ground truth boxes used to determine training labels batch_size_per_image int number of proposals to sample when training positive_fraction float target fraction of sampled proposals that should be positive images imagelist class imagelist instance representing n input images pred_objectness_logits list tensor a list of l elements element i is a tensor of shape n a hi w pred_anchor_deltas list tensor a list of l elements element i is a tensor of shape n a 4 hi wi anchors list torch tensor nested list of boxes anchors i j at n l stores anchor array for feature map l boundary_threshold int if 0 then anchors that extend beyond the image boundary by more than boundary_thresh are not used in training gt_boxes list boxes optional a list of n elements smooth_l1_beta float the transition point between l1 and l2 lossn when set to 0 the loss becomes l1 when inf it is ignored pred_anchor_deltas l n hi wi anchors n l 1 b here we loop over specific feature map not images append feature map proposals with shape n hi wi a b returns pred_objectness_logits list tensor n hi wi a reshape n a hi wi n hi wi a n hi wi a main classes this module is used in the original fpn to generate a downsampled p6 feature from p5 this module is used in retinanet to generate extra layers p6 and p7 from c5 feature use pad 1 instead of pad zero stride 2 conv stride 2 max pool the original msra resnet models have stride in the first 1x1 conv the subsequent fb torch resnet and caffe2 resne x t implementations have stride in the 3x3 conv some backbones require the input height and width to be divisible by a specific integer this is typically true for encoder decoder type networks with lateral connection e g fpn for which feature maps need to match dimension in the bottom up and top down paths set to 0 if no specific input size divisibility is required deprecated deprecated deprecated args stem nn module a stem module stages list list resnetblock several typically 4 stages each contains multiple class resnetblockbase num_classes none or int if none will not perform classification out_features list str name of the layers whose outputs should be returned in forward can be anything in stem linear or res2 if none will return the output of the last layer sec 5 1 in accurate large minibatch sgd training imagenet in 1 hour the 1000 way fully connected layer is initialized by drawing weights from a zero mean gaussian with std of 0 01 usually layers that produce the same feature map spatial size are defined as one stage under such definition stride_per_block 1 should all be 1 region of interest feature map pooler that supports pooling from one or more feature maps assumption that stride is a power of 2 a bunch of testing args feature_maps list torch tensor n c w h box_lists list torch tensor returns a tensor of shape n b channels output_size output_size only the pred boxes is the list of length n fun for each image too also i can experiment and do multiple images roiheads perform all per region computation in an r cnn it contains logic of cropping the regions extract per region features by the res 5 block in this case and make per region predictions res5 is 8x res2 self proposal_matcher matcher cfg roi_heads iou_thresholds cfg roi_heads iou_labels allow_low_quality_matches false modifications for vg in roi heads 1 change the stride of conv1 and shortcut in res5 block1 from 2 to 1 2 modifying all conv2 with padding 1 2 and dilation 1 2 res5 is 8x res2 see https github com airsplay py bottom up attention blob master detectron2 modeling roi_heads roi_heads py pooled to 1x1 for a set of image sizes and feature maps computes a set of anchors sizes list list int sizes i is the list of anchor sizes for feat map i 1 given in absolute lengths in units of the input image 2 they do not dynamically scale if the input image size changes aspect_ratios list list float strides list int stride of each input feature if one size or aspect ratio is specified and there are multiple feature maps then we broadcast anchors of that single size or aspect ratio returns list int each int is the number of anchors at every pixel location on that feature map anchors are continuous geometric rectangles centered on one feature map point sample we can later build the set of anchors for the entire feature map by tiling these tensors args features list torch tensor list of feature maps on which to generate anchors returns torch tensor a list of image elements rpn classification and regression heads uses a 3x3 conv to produce a shared hidden state from which one 1x1 conv predicts objectness logits for each anchor and a second 1x1 conv predicts bounding box deltas specifying how to deform each anchor into an object proposal standard rpn is shared across levels modifications for vg in rpn modeling proposal_generator rpn py use hidden dim instead fo the same dim as res4 in_channels 3x3 conv for the hidden representation 1x1 conv for predicting objectness logits 1x1 conv for predicting box2box transform deltas args features list tensor list of feature maps region proposal network introduced by the faster r cnn paper args images torch tensor input images of length n features dict str tensor gt_instances features is dict key block level v feature_map for rpn only models the proposals are the final output two linear layers for predicting fast r cnn outputs 1 proposal to detection box regression deltas 2 classification scores args input_size int channels or channels height width num_classes int cls_agnostic_bbox_reg bool box_dim int do 1 for background class modifications for vg in roi heads embedding num_classes 1 input_size 8 linear input_size input_size 8 input_size 4 linear input_size 4 num_attrs 1 b c b b b 256 b 2048 b 256 b 2304 load config if we don t provide a configuration try load model load from a pytorch checkpoint load from url or cache if already cached instantiate model convert old format to new format if needed from a pytorch state_dict copy state_dict so _load_from_state_dict can modify it set model in evaluation mode to deactivate dropout modules by default kwargs max_detections int return_tensors np pt none padding none max_detections pad_value int location cuda cpu run images through backbone generate proposals if none are available pool object features from either gt_boxes or from proposals prepare frcnn outputs and select top proposals will we pad
import itertools
import math
import os
from abc import ABCMeta, abstractmethod
from collections import OrderedDict, namedtuple
from typing import Dict, List, Tuple

import numpy as np
import torch
from torch import nn
from torch.nn.modules.batchnorm import BatchNorm2d
from torchvision.ops import RoIPool
from torchvision.ops.boxes import batched_nms, nms

from utils import WEIGHTS_NAME, Config, cached_path, hf_bucket_url, is_remote_url, load_checkpoint



def norm_box(boxes, raw_sizes):
    if not isinstance(boxes, torch.Tensor):
        normalized_boxes = boxes.copy()
    else:
        normalized_boxes = boxes.clone()
    normalized_boxes[:, :, (0, 2)] /= raw_sizes[:, 1]
    normalized_boxes[:, :, (1, 3)] /= raw_sizes[:, 0]
    return normalized_boxes


def pad_list_tensors(
    list_tensors,
    preds_per_image,
    max_detections=None,
    return_tensors=None,
    padding=None,
    pad_value=0,
    location=None,
):
    
    if location is None:
        location = "cpu"
    assert return_tensors in {"pt", "np", None}
    assert padding in {"max_detections", "max_batch", None}
    new = []
    if padding is None:
        if return_tensors is None:
            return list_tensors
        elif return_tensors == "pt":
            if not isinstance(list_tensors, torch.Tensor):
                return torch.stack(list_tensors).to(location)
            else:
                return list_tensors.to(location)
        else:
            if not isinstance(list_tensors, list):
                return np.array(list_tensors.to(location))
            else:
                return list_tensors.to(location)
    if padding == "max_detections":
        assert max_detections is not None, "specify max number of detections per batch"
    elif padding == "max_batch":
        max_detections = max(preds_per_image)
    for i in range(len(list_tensors)):
        too_small = False
        tensor_i = list_tensors.pop(0)
        if tensor_i.ndim < 2:
            too_small = True
            tensor_i = tensor_i.unsqueeze(-1)
        assert isinstance(tensor_i, torch.Tensor)
        tensor_i = nn.functional.pad(
            input=tensor_i,
            pad=(0, 0, 0, max_detections - preds_per_image[i]),
            mode="constant",
            value=pad_value,
        )
        if too_small:
            tensor_i = tensor_i.squeeze(-1)
        if return_tensors is None:
            if location == "cpu":
                tensor_i = tensor_i.cpu()
            tensor_i = tensor_i.tolist()
        if return_tensors == "np":
            if location == "cpu":
                tensor_i = tensor_i.cpu()
            tensor_i = tensor_i.numpy()
        else:
            if location == "cpu":
                tensor_i = tensor_i.cpu()
        new.append(tensor_i)
    if return_tensors == "np":
        return np.stack(new, axis=0)
    elif return_tensors == "pt" and not isinstance(new, torch.Tensor):
        return torch.stack(new, dim=0)
    else:
        return list_tensors


def do_nms(boxes, scores, image_shape, score_thresh, nms_thresh, mind, maxd):
    scores = scores[:, :-1]
    num_bbox_reg_classes = boxes.shape[1] // 4
    
    boxes = boxes.reshape(-1, 4)
    _clip_box(boxes, image_shape)
    boxes = boxes.view(-1, num_bbox_reg_classes, 4)  

    
    max_scores, max_classes = scores.max(1)  
    num_objs = boxes.size(0)
    boxes = boxes.view(-1, 4)
    idxs = torch.arange(num_objs).to(boxes.device) * num_bbox_reg_classes + max_classes
    max_boxes = boxes[idxs]  

    
    keep = nms(max_boxes, max_scores, nms_thresh)
    keep = keep[:maxd]
    if keep.shape[-1] >= mind and keep.shape[-1] <= maxd:
        max_boxes, max_scores = max_boxes[keep], max_scores[keep]
        classes = max_classes[keep]
        return max_boxes, max_scores, classes, keep
    else:
        return None



def _clip_box(tensor, box_size: Tuple[int, int]):
    assert torch.isfinite(tensor).all(), "Box tensor contains infinite or NaN!"
    h, w = box_size
    tensor[:, 0].clamp_(min=0, max=w)
    tensor[:, 1].clamp_(min=0, max=h)
    tensor[:, 2].clamp_(min=0, max=w)
    tensor[:, 3].clamp_(min=0, max=h)


def _nonempty_boxes(box, threshold: float = 0.0) -> torch.Tensor:
    widths = box[:, 2] - box[:, 0]
    heights = box[:, 3] - box[:, 1]
    keep = (widths > threshold) & (heights > threshold)
    return keep


def get_norm(norm, out_channels):
    if isinstance(norm, str):
        if len(norm) == 0:
            return None
        norm = {
            "BN": BatchNorm2d,
            "GN": lambda channels: nn.GroupNorm(32, channels),
            "nnSyncBN": nn.SyncBatchNorm,  
            "": lambda x: x,
        }[norm]
    return norm(out_channels)


def _create_grid_offsets(size: List[int], stride: int, offset: float, device):
    grid_height, grid_width = size
    shifts_x = torch.arange(
        offset * stride,
        grid_width * stride,
        step=stride,
        dtype=torch.float32,
        device=device,
    )
    shifts_y = torch.arange(
        offset * stride,
        grid_height * stride,
        step=stride,
        dtype=torch.float32,
        device=device,
    )

    shift_y, shift_x = torch.meshgrid(shifts_y, shifts_x)
    shift_x = shift_x.reshape(-1)
    shift_y = shift_y.reshape(-1)
    return shift_x, shift_y


def build_backbone(cfg):
    input_shape = ShapeSpec(channels=len(cfg.MODEL.PIXEL_MEAN))
    norm = cfg.RESNETS.NORM
    stem = BasicStem(
        in_channels=input_shape.channels,
        out_channels=cfg.RESNETS.STEM_OUT_CHANNELS,
        norm=norm,
        caffe_maxpool=cfg.MODEL.MAX_POOL,
    )
    freeze_at = cfg.BACKBONE.FREEZE_AT

    if freeze_at >= 1:
        for p in stem.parameters():
            p.requires_grad = False

    out_features = cfg.RESNETS.OUT_FEATURES
    depth = cfg.RESNETS.DEPTH
    num_groups = cfg.RESNETS.NUM_GROUPS
    width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
    bottleneck_channels = num_groups * width_per_group
    in_channels = cfg.RESNETS.STEM_OUT_CHANNELS
    out_channels = cfg.RESNETS.RES2_OUT_CHANNELS
    stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
    res5_dilation = cfg.RESNETS.RES5_DILATION
    assert res5_dilation in {1, 2}, "res5_dilation cannot be {}.".format(res5_dilation)

    num_blocks_per_stage = {50: [3, 4, 6, 3], 101: [3, 4, 23, 3], 152: [3, 8, 36, 3]}[depth]

    stages = []
    out_stage_idx = [{"res2": 2, "res3": 3, "res4": 4, "res5": 5}[f] for f in out_features]
    max_stage_idx = max(out_stage_idx)
    for idx, stage_idx in enumerate(range(2, max_stage_idx + 1)):
        dilation = res5_dilation if stage_idx == 5 else 1
        first_stride = 1 if idx == 0 or (stage_idx == 5 and dilation == 2) else 2
        stage_kargs = {
            "num_blocks": num_blocks_per_stage[idx],
            "first_stride": first_stride,
            "in_channels": in_channels,
            "bottleneck_channels": bottleneck_channels,
            "out_channels": out_channels,
            "num_groups": num_groups,
            "norm": norm,
            "stride_in_1x1": stride_in_1x1,
            "dilation": dilation,
        }

        stage_kargs["block_class"] = BottleneckBlock
        blocks = ResNet.make_stage(**stage_kargs)
        in_channels = out_channels
        out_channels *= 2
        bottleneck_channels *= 2

        if freeze_at >= stage_idx:
            for block in blocks:
                block.freeze()
        stages.append(blocks)

    return ResNet(stem, stages, out_features=out_features)


def find_top_rpn_proposals(
    proposals,
    pred_objectness_logits,
    images,
    image_sizes,
    nms_thresh,
    pre_nms_topk,
    post_nms_topk,
    min_box_side_len,
    training,
):
    
    num_images = len(images)
    device = proposals[0].device

    
    topk_scores = []  
    topk_proposals = []
    level_ids = []  
    batch_idx = torch.arange(num_images, device=device)
    for level_id, proposals_i, logits_i in zip(itertools.count(), proposals, pred_objectness_logits):
        Hi_Wi_A = logits_i.shape[1]
        num_proposals_i = min(pre_nms_topk, Hi_Wi_A)

        
        
        logits_i, idx = logits_i.sort(descending=True, dim=1)
        topk_scores_i = logits_i[batch_idx, :num_proposals_i]
        topk_idx = idx[batch_idx, :num_proposals_i]

        
        topk_proposals_i = proposals_i[batch_idx[:, None], topk_idx]  

        topk_proposals.append(topk_proposals_i)
        topk_scores.append(topk_scores_i)
        level_ids.append(torch.full((num_proposals_i,), level_id, dtype=torch.int64, device=device))

    
    topk_scores = torch.cat(topk_scores, dim=1)
    topk_proposals = torch.cat(topk_proposals, dim=1)
    level_ids = torch.cat(level_ids, dim=0)

    
    
    results = []
    for n, image_size in enumerate(image_sizes):
        boxes = topk_proposals[n]
        scores_per_img = topk_scores[n]
        
        _clip_box(boxes, image_size)
        
        keep = _nonempty_boxes(boxes, threshold=min_box_side_len)
        lvl = level_ids
        if keep.sum().item() != len(boxes):
            boxes, scores_per_img, lvl = (
                boxes[keep],
                scores_per_img[keep],
                level_ids[keep],
            )

        keep = batched_nms(boxes, scores_per_img, lvl, nms_thresh)
        keep = keep[:post_nms_topk]

        res = (boxes[keep], scores_per_img[keep])
        results.append(res)

    
    return results


def subsample_labels(labels, num_samples, positive_fraction, bg_label):
    
    positive = torch.nonzero((labels != -1) & (labels != bg_label)).squeeze(1)
    negative = torch.nonzero(labels == bg_label).squeeze(1)

    num_pos = int(num_samples * positive_fraction)
    
    num_pos = min(positive.numel(), num_pos)
    num_neg = num_samples - num_pos
    
    num_neg = min(negative.numel(), num_neg)

    
    perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]
    perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]

    pos_idx = positive[perm1]
    neg_idx = negative[perm2]
    return pos_idx, neg_idx


def add_ground_truth_to_proposals(gt_boxes, proposals):
    raise NotImplementedError()


def add_ground_truth_to_proposals_single_image(gt_boxes, proposals):
    raise NotImplementedError()


def _fmt_box_list(box_tensor, batch_index: int):
    repeated_index = torch.full(
        (len(box_tensor), 1),
        batch_index,
        dtype=box_tensor.dtype,
        device=box_tensor.device,
    )
    return torch.cat((repeated_index, box_tensor), dim=1)


def convert_boxes_to_pooler_format(box_lists: List[torch.Tensor]):
    pooler_fmt_boxes = torch.cat(
        [_fmt_box_list(box_list, i) for i, box_list in enumerate(box_lists)],
        dim=0,
    )
    return pooler_fmt_boxes


def assign_boxes_to_levels(
    box_lists: List[torch.Tensor],
    min_level: int,
    max_level: int,
    canonical_box_size: int,
    canonical_level: int,
):
    box_sizes = torch.sqrt(torch.cat([boxes.area() for boxes in box_lists]))
    
    level_assignments = torch.floor(canonical_level + torch.log2(box_sizes / canonical_box_size + 1e-8))
    
    
    level_assignments = torch.clamp(level_assignments, min=min_level, max=max_level)
    return level_assignments.to(torch.int64) - min_level



class _NewEmptyTensorOp(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x, new_shape):
        ctx.shape = x.shape
        return x.new_empty(new_shape)

    @staticmethod
    def backward(ctx, grad):
        shape = ctx.shape
        return _NewEmptyTensorOp.apply(grad, shape), None


class ShapeSpec(namedtuple("_ShapeSpec", ["channels", "height", "width", "stride"])):
    def __new__(cls, *, channels=None, height=None, width=None, stride=None):
        return super().__new__(cls, channels, height, width, stride)


class Box2BoxTransform(object):
    

    def __init__(self, weights: Tuple[float, float, float, float], scale_clamp: float = None):
        
        self.weights = weights
        if scale_clamp is not None:
            self.scale_clamp = scale_clamp
        else:
            
            self.scale_clamp = math.log(1000.0 / 16)

    def get_deltas(self, src_boxes, target_boxes):
        
        assert isinstance(src_boxes, torch.Tensor), type(src_boxes)
        assert isinstance(target_boxes, torch.Tensor), type(target_boxes)

        src_widths = src_boxes[:, 2] - src_boxes[:, 0]
        src_heights = src_boxes[:, 3] - src_boxes[:, 1]
        src_ctr_x = src_boxes[:, 0] + 0.5 * src_widths
        src_ctr_y = src_boxes[:, 1] + 0.5 * src_heights

        target_widths = target_boxes[:, 2] - target_boxes[:, 0]
        target_heights = target_boxes[:, 3] - target_boxes[:, 1]
        target_ctr_x = target_boxes[:, 0] + 0.5 * target_widths
        target_ctr_y = target_boxes[:, 1] + 0.5 * target_heights

        wx, wy, ww, wh = self.weights
        dx = wx * (target_ctr_x - src_ctr_x) / src_widths
        dy = wy * (target_ctr_y - src_ctr_y) / src_heights
        dw = ww * torch.log(target_widths / src_widths)
        dh = wh * torch.log(target_heights / src_heights)

        deltas = torch.stack((dx, dy, dw, dh), dim=1)
        assert (src_widths > 0).all().item(), "Input boxes to Box2BoxTransform are not valid!"
        return deltas

    def apply_deltas(self, deltas, boxes):
        
        boxes = boxes.to(deltas.dtype)

        widths = boxes[:, 2] - boxes[:, 0]
        heights = boxes[:, 3] - boxes[:, 1]
        ctr_x = boxes[:, 0] + 0.5 * widths
        ctr_y = boxes[:, 1] + 0.5 * heights

        wx, wy, ww, wh = self.weights
        dx = deltas[:, 0::4] / wx
        dy = deltas[:, 1::4] / wy
        dw = deltas[:, 2::4] / ww
        dh = deltas[:, 3::4] / wh

        
        dw = torch.clamp(dw, max=self.scale_clamp)
        dh = torch.clamp(dh, max=self.scale_clamp)

        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
        pred_w = torch.exp(dw) * widths[:, None]
        pred_h = torch.exp(dh) * heights[:, None]

        pred_boxes = torch.zeros_like(deltas)
        pred_boxes[:, 0::4] = pred_ctr_x - 0.5 * pred_w  
        pred_boxes[:, 1::4] = pred_ctr_y - 0.5 * pred_h  
        pred_boxes[:, 2::4] = pred_ctr_x + 0.5 * pred_w  
        pred_boxes[:, 3::4] = pred_ctr_y + 0.5 * pred_h  
        return pred_boxes


class Matcher(object):
    

    def __init__(
        self,
        thresholds: List[float],
        labels: List[int],
        allow_low_quality_matches: bool = False,
    ):
        
        thresholds = thresholds[:]
        assert thresholds[0] > 0
        thresholds.insert(0, -float("inf"))
        thresholds.append(float("inf"))
        assert all(low <= high for (low, high) in zip(thresholds[:-1], thresholds[1:]))
        assert all(label_i in [-1, 0, 1] for label_i in labels)
        assert len(labels) == len(thresholds) - 1
        self.thresholds = thresholds
        self.labels = labels
        self.allow_low_quality_matches = allow_low_quality_matches

    def __call__(self, match_quality_matrix):
        
        assert match_quality_matrix.dim() == 2
        if match_quality_matrix.numel() == 0:
            default_matches = match_quality_matrix.new_full((match_quality_matrix.size(1),), 0, dtype=torch.int64)
            
            
            
            
            default_match_labels = match_quality_matrix.new_full(
                (match_quality_matrix.size(1),), self.labels[0], dtype=torch.int8
            )
            return default_matches, default_match_labels

        assert torch.all(match_quality_matrix >= 0)

        
        
        matched_vals, matches = match_quality_matrix.max(dim=0)

        match_labels = matches.new_full(matches.size(), 1, dtype=torch.int8)

        for l, low, high in zip(self.labels, self.thresholds[:-1], self.thresholds[1:]):
            low_high = (matched_vals >= low) & (matched_vals < high)
            match_labels[low_high] = l

        if self.allow_low_quality_matches:
            self.set_low_quality_matches_(match_labels, match_quality_matrix)

        return matches, match_labels

    def set_low_quality_matches_(self, match_labels, match_quality_matrix):
        
        
        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
        
        
        
        of_quality_inds = match_quality_matrix == highest_quality_foreach_gt[:, None]
        if of_quality_inds.dim() == 0:
            (_, pred_inds_with_highest_quality) = of_quality_inds.unsqueeze(0).nonzero().unbind(1)
        else:
            (_, pred_inds_with_highest_quality) = of_quality_inds.nonzero().unbind(1)
        match_labels[pred_inds_with_highest_quality] = 1


class RPNOutputs(object):
    def __init__(
        self,
        box2box_transform,
        anchor_matcher,
        batch_size_per_image,
        positive_fraction,
        images,
        pred_objectness_logits,
        pred_anchor_deltas,
        anchors,
        boundary_threshold=0,
        gt_boxes=None,
        smooth_l1_beta=0.0,
    ):
        
        self.box2box_transform = box2box_transform
        self.anchor_matcher = anchor_matcher
        self.batch_size_per_image = batch_size_per_image
        self.positive_fraction = positive_fraction
        self.pred_objectness_logits = pred_objectness_logits
        self.pred_anchor_deltas = pred_anchor_deltas

        self.anchors = anchors
        self.gt_boxes = gt_boxes
        self.num_feature_maps = len(pred_objectness_logits)
        self.num_images = len(images)
        self.boundary_threshold = boundary_threshold
        self.smooth_l1_beta = smooth_l1_beta

    def _get_ground_truth(self):
        raise NotImplementedError()

    def predict_proposals(self):
        
        
        
        proposals = []
        anchors = self.anchors.transpose(0, 1)
        for anchors_i, pred_anchor_deltas_i in zip(anchors, self.pred_anchor_deltas):
            B = anchors_i.size(-1)
            N, _, Hi, Wi = pred_anchor_deltas_i.shape
            anchors_i = anchors_i.flatten(start_dim=0, end_dim=1)
            pred_anchor_deltas_i = pred_anchor_deltas_i.view(N, -1, B, Hi, Wi).permute(0, 3, 4, 1, 2).reshape(-1, B)
            proposals_i = self.box2box_transform.apply_deltas(pred_anchor_deltas_i, anchors_i)
            
            proposals.append(proposals_i.view(N, -1, B))
        proposals = torch.stack(proposals)
        return proposals

    def predict_objectness_logits(self):
        
        pred_objectness_logits = [
            
            score.permute(0, 2, 3, 1).reshape(self.num_images, -1)
            for score in self.pred_objectness_logits
        ]
        return pred_objectness_logits



class Conv2d(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        norm = kwargs.pop("norm", None)
        activation = kwargs.pop("activation", None)
        super().__init__(*args, **kwargs)

        self.norm = norm
        self.activation = activation

    def forward(self, x):
        if x.numel() == 0 and self.training:
            assert not isinstance(self.norm, nn.SyncBatchNorm)
        if x.numel() == 0:
            assert not isinstance(self.norm, nn.GroupNorm)
            output_shape = [
                (i + 2 * p - (di * (k - 1) + 1)) // s + 1
                for i, p, di, k, s in zip(
                    x.shape[-2:],
                    self.padding,
                    self.dilation,
                    self.kernel_size,
                    self.stride,
                )
            ]
            output_shape = [x.shape[0], self.weight.shape[0]] + output_shape
            empty = _NewEmptyTensorOp.apply(x, output_shape)
            if self.training:
                _dummy = sum(x.view(-1)[0] for x in self.parameters()) * 0.0
                return empty + _dummy
            else:
                return empty

        x = super().forward(x)
        if self.norm is not None:
            x = self.norm(x)
        if self.activation is not None:
            x = self.activation(x)
        return x


class LastLevelMaxPool(nn.Module):
    

    def __init__(self):
        super().__init__()
        self.num_levels = 1
        self.in_feature = "p5"

    def forward(self, x):
        return [nn.functional.max_pool2d(x, kernel_size=1, stride=2, padding=0)]


class LastLevelP6P7(nn.Module):
    

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.num_levels = 2
        self.in_feature = "res5"
        self.p6 = nn.Conv2d(in_channels, out_channels, 3, 2, 1)
        self.p7 = nn.Conv2d(out_channels, out_channels, 3, 2, 1)

    def forward(self, c5):
        p6 = self.p6(c5)
        p7 = self.p7(nn.functional.relu(p6))
        return [p6, p7]


class BasicStem(nn.Module):
    def __init__(self, in_channels=3, out_channels=64, norm="BN", caffe_maxpool=False):
        super().__init__()
        self.conv1 = Conv2d(
            in_channels,
            out_channels,
            kernel_size=7,
            stride=2,
            padding=3,
            bias=False,
            norm=get_norm(norm, out_channels),
        )
        self.caffe_maxpool = caffe_maxpool
        

    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu_(x)
        if self.caffe_maxpool:
            x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=0, ceil_mode=True)
        else:
            x = nn.functional.max_pool2d(x, kernel_size=3, stride=2, padding=1)
        return x

    @property
    def out_channels(self):
        return self.conv1.out_channels

    @property
    def stride(self):
        return 4  


class ResNetBlockBase(nn.Module):
    def __init__(self, in_channels, out_channels, stride):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.stride = stride

    def freeze(self):
        for p in self.parameters():
            p.requires_grad = False
        return self


class BottleneckBlock(ResNetBlockBase):
    def __init__(
        self,
        in_channels,
        out_channels,
        bottleneck_channels,
        stride=1,
        num_groups=1,
        norm="BN",
        stride_in_1x1=False,
        dilation=1,
    ):
        super().__init__(in_channels, out_channels, stride)

        if in_channels != out_channels:
            self.shortcut = Conv2d(
                in_channels,
                out_channels,
                kernel_size=1,
                stride=stride,
                bias=False,
                norm=get_norm(norm, out_channels),
            )
        else:
            self.shortcut = None

        
        
        
        stride_1x1, stride_3x3 = (stride, 1) if stride_in_1x1 else (1, stride)

        self.conv1 = Conv2d(
            in_channels,
            bottleneck_channels,
            kernel_size=1,
            stride=stride_1x1,
            bias=False,
            norm=get_norm(norm, bottleneck_channels),
        )

        self.conv2 = Conv2d(
            bottleneck_channels,
            bottleneck_channels,
            kernel_size=3,
            stride=stride_3x3,
            padding=1 * dilation,
            bias=False,
            groups=num_groups,
            dilation=dilation,
            norm=get_norm(norm, bottleneck_channels),
        )

        self.conv3 = Conv2d(
            bottleneck_channels,
            out_channels,
            kernel_size=1,
            bias=False,
            norm=get_norm(norm, out_channels),
        )

    def forward(self, x):
        out = self.conv1(x)
        out = nn.functional.relu_(out)

        out = self.conv2(out)
        out = nn.functional.relu_(out)

        out = self.conv3(out)

        if self.shortcut is not None:
            shortcut = self.shortcut(x)
        else:
            shortcut = x

        out += shortcut
        out = nn.functional.relu_(out)
        return out


class Backbone(nn.Module, metaclass=ABCMeta):
    def __init__(self):
        super().__init__()

    @abstractmethod
    def forward(self):
        pass

    @property
    def size_divisibility(self):
        
        return 0

    def output_shape(self):
        return {
            name: ShapeSpec(
                channels=self._out_feature_channels[name],
                stride=self._out_feature_strides[name],
            )
            for name in self._out_features
        }

    @property
    def out_features(self):
        
        return self._out_features

    @property
    def out_feature_strides(self):
        
        return {f: self._out_feature_strides[f] for f in self._out_features}

    @property
    def out_feature_channels(self):
        
        return {f: self._out_feature_channels[f] for f in self._out_features}


class ResNet(Backbone):
    def __init__(self, stem, stages, num_classes=None, out_features=None):
        
        super(ResNet, self).__init__()
        self.stem = stem
        self.num_classes = num_classes

        current_stride = self.stem.stride
        self._out_feature_strides = {"stem": current_stride}
        self._out_feature_channels = {"stem": self.stem.out_channels}

        self.stages_and_names = []
        for i, blocks in enumerate(stages):
            for block in blocks:
                assert isinstance(block, ResNetBlockBase), block
                curr_channels = block.out_channels
            stage = nn.Sequential(*blocks)
            name = "res" + str(i + 2)
            self.add_module(name, stage)
            self.stages_and_names.append((stage, name))
            self._out_feature_strides[name] = current_stride = int(
                current_stride * np.prod([k.stride for k in blocks])
            )
            self._out_feature_channels[name] = blocks[-1].out_channels

        if num_classes is not None:
            self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
            self.linear = nn.Linear(curr_channels, num_classes)

            
            
            
            nn.init.normal_(self.linear.weight, stddev=0.01)
            name = "linear"

        if out_features is None:
            out_features = [name]
        self._out_features = out_features
        assert len(self._out_features)
        children = [x[0] for x in self.named_children()]
        for out_feature in self._out_features:
            assert out_feature in children, "Available children: {}".format(", ".join(children))

    def forward(self, x):
        outputs = {}
        x = self.stem(x)
        if "stem" in self._out_features:
            outputs["stem"] = x
        for stage, name in self.stages_and_names:
            x = stage(x)
            if name in self._out_features:
                outputs[name] = x
        if self.num_classes is not None:
            x = self.avgpool(x)
            x = self.linear(x)
            if "linear" in self._out_features:
                outputs["linear"] = x
        return outputs

    def output_shape(self):
        return {
            name: ShapeSpec(
                channels=self._out_feature_channels[name],
                stride=self._out_feature_strides[name],
            )
            for name in self._out_features
        }

    @staticmethod
    def make_stage(
        block_class,
        num_blocks,
        first_stride=None,
        *,
        in_channels,
        out_channels,
        **kwargs,
    ):
        
        if first_stride is not None:
            assert "stride" not in kwargs and "stride_per_block" not in kwargs
            kwargs["stride_per_block"] = [first_stride] + [1] * (num_blocks - 1)
        blocks = []
        for i in range(num_blocks):
            curr_kwargs = {}
            for k, v in kwargs.items():
                if k.endswith("_per_block"):
                    assert (
                        len(v) == num_blocks
                    ), f"Argument '{k}' of make_stage should have the same length as num_blocks={num_blocks}."
                    newk = k[: -len("_per_block")]
                    assert newk not in kwargs, f"Cannot call make_stage with both {k} and {newk}!"
                    curr_kwargs[newk] = v[i]
                else:
                    curr_kwargs[k] = v

            blocks.append(block_class(in_channels=in_channels, out_channels=out_channels, **curr_kwargs))
            in_channels = out_channels

        return blocks


class ROIPooler(nn.Module):
    

    def __init__(
        self,
        output_size,
        scales,
        sampling_ratio,
        canonical_box_size=224,
        canonical_level=4,
    ):
        super().__init__()
        
        min_level = -math.log2(scales[0])
        max_level = -math.log2(scales[-1])

        
        assert math.isclose(min_level, int(min_level)) and math.isclose(max_level, int(max_level))
        assert len(scales) == max_level - min_level + 1, "not pyramid"
        assert 0 < min_level and min_level <= max_level
        if isinstance(output_size, int):
            output_size = (output_size, output_size)
        assert len(output_size) == 2 and isinstance(output_size[0], int) and isinstance(output_size[1], int)
        if len(scales) > 1:
            assert min_level <= canonical_level and canonical_level <= max_level
        assert canonical_box_size > 0

        self.output_size = output_size
        self.min_level = int(min_level)
        self.max_level = int(max_level)
        self.level_poolers = nn.ModuleList(RoIPool(output_size, spatial_scale=scale) for scale in scales)
        self.canonical_level = canonical_level
        self.canonical_box_size = canonical_box_size

    def forward(self, feature_maps, boxes):
        
        x = list(feature_maps.values())
        num_level_assignments = len(self.level_poolers)
        assert len(x) == num_level_assignments and len(boxes) == x[0].size(0)

        pooler_fmt_boxes = convert_boxes_to_pooler_format(boxes)

        if num_level_assignments == 1:
            return self.level_poolers[0](x[0], pooler_fmt_boxes)

        level_assignments = assign_boxes_to_levels(
            boxes,
            self.min_level,
            self.max_level,
            self.canonical_box_size,
            self.canonical_level,
        )

        num_boxes = len(pooler_fmt_boxes)
        num_channels = x[0].shape[1]
        output_size = self.output_size[0]

        dtype, device = x[0].dtype, x[0].device
        output = torch.zeros(
            (num_boxes, num_channels, output_size, output_size),
            dtype=dtype,
            device=device,
        )

        for level, (x_level, pooler) in enumerate(zip(x, self.level_poolers)):
            inds = torch.nonzero(level_assignments == level).squeeze(1)
            pooler_fmt_boxes_level = pooler_fmt_boxes[inds]
            output[inds] = pooler(x_level, pooler_fmt_boxes_level)

        return output


class ROIOutputs(object):
    def __init__(self, cfg, training=False):
        self.smooth_l1_beta = cfg.ROI_BOX_HEAD.SMOOTH_L1_BETA
        self.box2box_transform = Box2BoxTransform(weights=cfg.ROI_BOX_HEAD.BBOX_REG_WEIGHTS)
        self.training = training
        self.score_thresh = cfg.ROI_HEADS.SCORE_THRESH_TEST
        self.min_detections = cfg.MIN_DETECTIONS
        self.max_detections = cfg.MAX_DETECTIONS

        nms_thresh = cfg.ROI_HEADS.NMS_THRESH_TEST
        if not isinstance(nms_thresh, list):
            nms_thresh = [nms_thresh]
        self.nms_thresh = nms_thresh

    def _predict_boxes(self, proposals, box_deltas, preds_per_image):
        num_pred = box_deltas.size(0)
        B = proposals[0].size(-1)
        K = box_deltas.size(-1) // B
        box_deltas = box_deltas.view(num_pred * K, B)
        proposals = torch.cat(proposals, dim=0).unsqueeze(-2).expand(num_pred, K, B)
        proposals = proposals.reshape(-1, B)
        boxes = self.box2box_transform.apply_deltas(box_deltas, proposals)
        return boxes.view(num_pred, K * B).split(preds_per_image, dim=0)

    def _predict_objs(self, obj_logits, preds_per_image):
        probs = nn.functional.softmax(obj_logits, dim=-1)
        probs = probs.split(preds_per_image, dim=0)
        return probs

    def _predict_attrs(self, attr_logits, preds_per_image):
        attr_logits = attr_logits[..., :-1].softmax(-1)
        attr_probs, attrs = attr_logits.max(-1)
        return attr_probs.split(preds_per_image, dim=0), attrs.split(preds_per_image, dim=0)

    @torch.no_grad()
    def inference(
        self,
        obj_logits,
        attr_logits,
        box_deltas,
        pred_boxes,
        features,
        sizes,
        scales=None,
    ):
        
        preds_per_image = [p.size(0) for p in pred_boxes]
        boxes_all = self._predict_boxes(pred_boxes, box_deltas, preds_per_image)
        obj_scores_all = self._predict_objs(obj_logits, preds_per_image)  
        attr_probs_all, attrs_all = self._predict_attrs(attr_logits, preds_per_image)
        features = features.split(preds_per_image, dim=0)

        
        final_results = []
        zipped = zip(boxes_all, obj_scores_all, attr_probs_all, attrs_all, sizes)
        for i, (boxes, obj_scores, attr_probs, attrs, size) in enumerate(zipped):
            for nms_t in self.nms_thresh:
                outputs = do_nms(
                    boxes,
                    obj_scores,
                    size,
                    self.score_thresh,
                    nms_t,
                    self.min_detections,
                    self.max_detections,
                )
                if outputs is not None:
                    max_boxes, max_scores, classes, ids = outputs
                    break

            if scales is not None:
                scale_yx = scales[i]
                max_boxes[:, 0::2] *= scale_yx[1]
                max_boxes[:, 1::2] *= scale_yx[0]

            final_results.append(
                (
                    max_boxes,
                    classes,
                    max_scores,
                    attrs[ids],
                    attr_probs[ids],
                    features[i][ids],
                )
            )
        boxes, classes, class_probs, attrs, attr_probs, roi_features = map(list, zip(*final_results))
        return boxes, classes, class_probs, attrs, attr_probs, roi_features

    def training(self, obj_logits, attr_logits, box_deltas, pred_boxes, features, sizes):
        pass

    def __call__(
        self,
        obj_logits,
        attr_logits,
        box_deltas,
        pred_boxes,
        features,
        sizes,
        scales=None,
    ):
        if self.training:
            raise NotImplementedError()
        return self.inference(
            obj_logits,
            attr_logits,
            box_deltas,
            pred_boxes,
            features,
            sizes,
            scales=scales,
        )


class Res5ROIHeads(nn.Module):
    

    def __init__(self, cfg, input_shape):
        super().__init__()
        self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
        self.positive_sample_fraction = cfg.ROI_HEADS.POSITIVE_FRACTION
        self.in_features = cfg.ROI_HEADS.IN_FEATURES
        self.num_classes = cfg.ROI_HEADS.NUM_CLASSES
        self.proposal_append_gt = cfg.ROI_HEADS.PROPOSAL_APPEND_GT
        self.feature_strides = {k: v.stride for k, v in input_shape.items()}
        self.feature_channels = {k: v.channels for k, v in input_shape.items()}
        self.cls_agnostic_bbox_reg = cfg.ROI_BOX_HEAD.CLS_AGNOSTIC_BBOX_REG
        self.stage_channel_factor = 2**3  
        self.out_channels = cfg.RESNETS.RES2_OUT_CHANNELS * self.stage_channel_factor

        
        
        
        
        

        pooler_resolution = cfg.ROI_BOX_HEAD.POOLER_RESOLUTION
        pooler_scales = (1.0 / self.feature_strides[self.in_features[0]],)
        sampling_ratio = cfg.ROI_BOX_HEAD.POOLER_SAMPLING_RATIO
        res5_halve = cfg.ROI_BOX_HEAD.RES5HALVE
        use_attr = cfg.ROI_BOX_HEAD.ATTR
        num_attrs = cfg.ROI_BOX_HEAD.NUM_ATTRS

        self.pooler = ROIPooler(
            output_size=pooler_resolution,
            scales=pooler_scales,
            sampling_ratio=sampling_ratio,
        )

        self.res5 = self._build_res5_block(cfg)
        if not res5_halve:
            
            self.res5[0].conv1.stride = (1, 1)
            self.res5[0].shortcut.stride = (1, 1)
            for i in range(3):
                self.res5[i].conv2.padding = (2, 2)
                self.res5[i].conv2.dilation = (2, 2)

        self.box_predictor = FastRCNNOutputLayers(
            self.out_channels,
            self.num_classes,
            self.cls_agnostic_bbox_reg,
            use_attr=use_attr,
            num_attrs=num_attrs,
        )

    def _build_res5_block(self, cfg):
        stage_channel_factor = self.stage_channel_factor  
        num_groups = cfg.RESNETS.NUM_GROUPS
        width_per_group = cfg.RESNETS.WIDTH_PER_GROUP
        bottleneck_channels = num_groups * width_per_group * stage_channel_factor
        out_channels = self.out_channels
        stride_in_1x1 = cfg.RESNETS.STRIDE_IN_1X1
        norm = cfg.RESNETS.NORM

        blocks = ResNet.make_stage(
            BottleneckBlock,
            3,
            first_stride=2,
            in_channels=out_channels // 2,
            bottleneck_channels=bottleneck_channels,
            out_channels=out_channels,
            num_groups=num_groups,
            norm=norm,
            stride_in_1x1=stride_in_1x1,
        )
        return nn.Sequential(*blocks)

    def _shared_roi_transform(self, features, boxes):
        x = self.pooler(features, boxes)
        return self.res5(x)

    def forward(self, features, proposal_boxes, gt_boxes=None):
        if self.training:
            
            raise NotImplementedError()

        assert not proposal_boxes[0].requires_grad
        box_features = self._shared_roi_transform(features, proposal_boxes)
        feature_pooled = box_features.mean(dim=[2, 3])  
        obj_logits, attr_logits, pred_proposal_deltas = self.box_predictor(feature_pooled)
        return obj_logits, attr_logits, pred_proposal_deltas, feature_pooled


class AnchorGenerator(nn.Module):
    

    def __init__(self, cfg, input_shape: List[ShapeSpec]):
        super().__init__()
        sizes = cfg.ANCHOR_GENERATOR.SIZES
        aspect_ratios = cfg.ANCHOR_GENERATOR.ASPECT_RATIOS
        self.strides = [x.stride for x in input_shape]
        self.offset = cfg.ANCHOR_GENERATOR.OFFSET
        assert 0.0 <= self.offset < 1.0, self.offset

        

        self.num_features = len(self.strides)
        self.cell_anchors = nn.ParameterList(self._calculate_anchors(sizes, aspect_ratios))
        self._spacial_feat_dim = 4

    def _calculate_anchors(self, sizes, aspect_ratios):
        
        
        if len(sizes) == 1:
            sizes *= self.num_features
        if len(aspect_ratios) == 1:
            aspect_ratios *= self.num_features
        assert self.num_features == len(sizes)
        assert self.num_features == len(aspect_ratios)

        cell_anchors = [self.generate_cell_anchors(s, a).float() for s, a in zip(sizes, aspect_ratios)]

        return cell_anchors

    @property
    def box_dim(self):
        return self._spacial_feat_dim

    @property
    def num_cell_anchors(self):
        
        return [len(cell_anchors) for cell_anchors in self.cell_anchors]

    def grid_anchors(self, grid_sizes):
        anchors = []
        for size, stride, base_anchors in zip(grid_sizes, self.strides, self.cell_anchors):
            shift_x, shift_y = _create_grid_offsets(size, stride, self.offset, base_anchors.device)
            shifts = torch.stack((shift_x, shift_y, shift_x, shift_y), dim=1)

            anchors.append((shifts.view(-1, 1, 4) + base_anchors.view(1, -1, 4)).reshape(-1, 4))

        return anchors

    def generate_cell_anchors(self, sizes=(32, 64, 128, 256, 512), aspect_ratios=(0.5, 1, 2)):
        

        anchors = []
        for size in sizes:
            area = size**2.0
            for aspect_ratio in aspect_ratios:
                w = math.sqrt(area / aspect_ratio)
                h = aspect_ratio * w
                x0, y0, x1, y1 = -w / 2.0, -h / 2.0, w / 2.0, h / 2.0
                anchors.append([x0, y0, x1, y1])
        return nn.Parameter(torch.tensor(anchors))

    def forward(self, features):
        
        num_images = features[0].size(0)
        grid_sizes = [feature_map.shape[-2:] for feature_map in features]
        anchors_over_all_feature_maps = self.grid_anchors(grid_sizes)
        anchors_over_all_feature_maps = torch.stack(anchors_over_all_feature_maps)
        return anchors_over_all_feature_maps.unsqueeze(0).repeat_interleave(num_images, dim=0)


class RPNHead(nn.Module):
    

    def __init__(self, cfg, input_shape: List[ShapeSpec]):
        super().__init__()

        
        in_channels = [s.channels for s in input_shape]
        assert len(set(in_channels)) == 1, "Each level must have the same channel!"
        in_channels = in_channels[0]

        anchor_generator = AnchorGenerator(cfg, input_shape)
        num_cell_anchors = anchor_generator.num_cell_anchors
        box_dim = anchor_generator.box_dim
        assert len(set(num_cell_anchors)) == 1, "Each level must have the same number of cell anchors"
        num_cell_anchors = num_cell_anchors[0]

        if cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS == -1:
            hid_channels = in_channels
        else:
            hid_channels = cfg.PROPOSAL_GENERATOR.HIDDEN_CHANNELS
            
            

        
        self.conv = nn.Conv2d(in_channels, hid_channels, kernel_size=3, stride=1, padding=1)
        
        self.objectness_logits = nn.Conv2d(hid_channels, num_cell_anchors, kernel_size=1, stride=1)
        
        self.anchor_deltas = nn.Conv2d(hid_channels, num_cell_anchors * box_dim, kernel_size=1, stride=1)

        for layer in [self.conv, self.objectness_logits, self.anchor_deltas]:
            nn.init.normal_(layer.weight, std=0.01)
            nn.init.constant_(layer.bias, 0)

    def forward(self, features):
        
        pred_objectness_logits = []
        pred_anchor_deltas = []
        for x in features:
            t = nn.functional.relu(self.conv(x))
            pred_objectness_logits.append(self.objectness_logits(t))
            pred_anchor_deltas.append(self.anchor_deltas(t))
        return pred_objectness_logits, pred_anchor_deltas


class RPN(nn.Module):
    

    def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]):
        super().__init__()

        self.min_box_side_len = cfg.PROPOSAL_GENERATOR.MIN_SIZE
        self.in_features = cfg.RPN.IN_FEATURES
        self.nms_thresh = cfg.RPN.NMS_THRESH
        self.batch_size_per_image = cfg.RPN.BATCH_SIZE_PER_IMAGE
        self.positive_fraction = cfg.RPN.POSITIVE_FRACTION
        self.smooth_l1_beta = cfg.RPN.SMOOTH_L1_BETA
        self.loss_weight = cfg.RPN.LOSS_WEIGHT

        self.pre_nms_topk = {
            True: cfg.RPN.PRE_NMS_TOPK_TRAIN,
            False: cfg.RPN.PRE_NMS_TOPK_TEST,
        }
        self.post_nms_topk = {
            True: cfg.RPN.POST_NMS_TOPK_TRAIN,
            False: cfg.RPN.POST_NMS_TOPK_TEST,
        }
        self.boundary_threshold = cfg.RPN.BOUNDARY_THRESH

        self.anchor_generator = AnchorGenerator(cfg, [input_shape[f] for f in self.in_features])
        self.box2box_transform = Box2BoxTransform(weights=cfg.RPN.BBOX_REG_WEIGHTS)
        self.anchor_matcher = Matcher(
            cfg.RPN.IOU_THRESHOLDS,
            cfg.RPN.IOU_LABELS,
            allow_low_quality_matches=True,
        )
        self.rpn_head = RPNHead(cfg, [input_shape[f] for f in self.in_features])

    def training(self, images, image_shapes, features, gt_boxes):
        pass

    def inference(self, outputs, images, image_shapes, features, gt_boxes=None):
        outputs = find_top_rpn_proposals(
            outputs.predict_proposals(),
            outputs.predict_objectness_logits(),
            images,
            image_shapes,
            self.nms_thresh,
            self.pre_nms_topk[self.training],
            self.post_nms_topk[self.training],
            self.min_box_side_len,
            self.training,
        )

        results = []
        for img in outputs:
            im_boxes, img_box_logits = img
            img_box_logits, inds = img_box_logits.sort(descending=True)
            im_boxes = im_boxes[inds]
            results.append((im_boxes, img_box_logits))

        (proposal_boxes, logits) = tuple(map(list, zip(*results)))
        return proposal_boxes, logits

    def forward(self, images, image_shapes, features, gt_boxes=None):
        
        
        features = [features[f] for f in self.in_features]
        pred_objectness_logits, pred_anchor_deltas = self.rpn_head(features)
        anchors = self.anchor_generator(features)
        outputs = RPNOutputs(
            self.box2box_transform,
            self.anchor_matcher,
            self.batch_size_per_image,
            self.positive_fraction,
            images,
            pred_objectness_logits,
            pred_anchor_deltas,
            anchors,
            self.boundary_threshold,
            gt_boxes,
            self.smooth_l1_beta,
        )
        

        if self.training:
            raise NotImplementedError()
            return self.training(outputs, images, image_shapes, features, gt_boxes)
        else:
            return self.inference(outputs, images, image_shapes, features, gt_boxes)


class FastRCNNOutputLayers(nn.Module):
    

    def __init__(
        self,
        input_size,
        num_classes,
        cls_agnostic_bbox_reg,
        box_dim=4,
        use_attr=False,
        num_attrs=-1,
    ):
        
        super().__init__()

        if not isinstance(input_size, int):
            input_size = np.prod(input_size)

        
        self.cls_score = nn.Linear(input_size, num_classes + 1)
        num_bbox_reg_classes = 1 if cls_agnostic_bbox_reg else num_classes
        self.bbox_pred = nn.Linear(input_size, num_bbox_reg_classes * box_dim)

        self.use_attr = use_attr
        if use_attr:
            
            self.cls_embedding = nn.Embedding(num_classes + 1, input_size // 8)
            self.fc_attr = nn.Linear(input_size + input_size // 8, input_size // 4)
            self.attr_score = nn.Linear(input_size // 4, num_attrs + 1)

        nn.init.normal_(self.cls_score.weight, std=0.01)
        nn.init.normal_(self.bbox_pred.weight, std=0.001)
        for item in [self.cls_score, self.bbox_pred]:
            nn.init.constant_(item.bias, 0)

    def forward(self, roi_features):
        if roi_features.dim() > 2:
            roi_features = torch.flatten(roi_features, start_dim=1)
        scores = self.cls_score(roi_features)
        proposal_deltas = self.bbox_pred(roi_features)
        if self.use_attr:
            _, max_class = scores.max(-1)  
            cls_emb = self.cls_embedding(max_class)  
            roi_features = torch.cat([roi_features, cls_emb], -1)  
            roi_features = self.fc_attr(roi_features)
            roi_features = nn.functional.relu(roi_features)
            attr_scores = self.attr_score(roi_features)
            return scores, attr_scores, proposal_deltas
        else:
            return scores, proposal_deltas


class GeneralizedRCNN(nn.Module):
    def __init__(self, cfg):
        super().__init__()

        self.device = torch.device(cfg.MODEL.DEVICE)
        self.backbone = build_backbone(cfg)
        self.proposal_generator = RPN(cfg, self.backbone.output_shape())
        self.roi_heads = Res5ROIHeads(cfg, self.backbone.output_shape())
        self.roi_outputs = ROIOutputs(cfg)
        self.to(self.device)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        state_dict = kwargs.pop("state_dict", None)
        cache_dir = kwargs.pop("cache_dir", None)
        from_tf = kwargs.pop("from_tf", False)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", False)
        use_cdn = kwargs.pop("use_cdn", True)

        
        if not isinstance(config, Config):
            config_path = config if config is not None else pretrained_model_name_or_path
            
            config = Config.from_pretrained(
                config_path,
                cache_dir=cache_dir,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
            )

        
        if pretrained_model_name_or_path is not None:
            if os.path.isdir(pretrained_model_name_or_path):
                if os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
                    
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
                else:
                    raise EnvironmentError(
                        "Error no file named {} found in directory {} ".format(
                            WEIGHTS_NAME,
                            pretrained_model_name_or_path,
                        )
                    )
            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
                archive_file = pretrained_model_name_or_path
            elif os.path.isfile(pretrained_model_name_or_path + ".index"):
                assert from_tf, "We found a TensorFlow checkpoint at {}, please set from_tf to True to load from this checkpoint".format(
                    pretrained_model_name_or_path + ".index"
                )
                archive_file = pretrained_model_name_or_path + ".index"
            else:
                archive_file = hf_bucket_url(
                    pretrained_model_name_or_path,
                    filename=WEIGHTS_NAME,
                    use_cdn=use_cdn,
                )

            try:
                
                resolved_archive_file = cached_path(
                    archive_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    local_files_only=local_files_only,
                )
                if resolved_archive_file is None:
                    raise EnvironmentError
            except EnvironmentError:
                msg = f"Can't load weights for '{pretrained_model_name_or_path}'."
                raise EnvironmentError(msg)

            if resolved_archive_file == archive_file:
                print("loading weights file {}".format(archive_file))
            else:
                print("loading weights file {} from cache at {}".format(archive_file, resolved_archive_file))
        else:
            resolved_archive_file = None

        
        model = cls(config)

        if state_dict is None:
            try:
                try:
                    state_dict = torch.load(resolved_archive_file, map_location="cpu")
                except Exception:
                    state_dict = load_checkpoint(resolved_archive_file)

            except Exception:
                raise OSError(
                    "Unable to load weights from pytorch checkpoint file. "
                    "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True. "
                )

        missing_keys = []
        unexpected_keys = []
        error_msgs = []

        
        old_keys = []
        new_keys = []
        for key in state_dict.keys():
            new_key = None
            if "gamma" in key:
                new_key = key.replace("gamma", "weight")
            if "beta" in key:
                new_key = key.replace("beta", "bias")
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)
        for old_key, new_key in zip(old_keys, new_keys):
            state_dict[new_key] = state_dict.pop(old_key)

        
        metadata = getattr(state_dict, "_metadata", None)
        state_dict = state_dict.copy()
        if metadata is not None:
            state_dict._metadata = metadata

        model_to_load = model
        model_to_load.load_state_dict(state_dict)

        if model.__class__.__name__ != model_to_load.__class__.__name__:
            base_model_state_dict = model_to_load.state_dict().keys()
            head_model_state_dict_without_base_prefix = [
                key.split(cls.base_model_prefix + ".")[-1] for key in model.state_dict().keys()
            ]
            missing_keys.extend(head_model_state_dict_without_base_prefix - base_model_state_dict)

        if len(unexpected_keys) > 0:
            print(
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when"
                f" initializing {model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are"
                f" initializing {model.__class__.__name__} from the checkpoint of a model trained on another task or"
                " with another architecture (e.g. initializing a BertForSequenceClassification model from a"
                " BertForPreTraining model).\n- This IS NOT expected if you are initializing"
                f" {model.__class__.__name__} from the checkpoint of a model that you expect to be exactly identical"
                " (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
            )
        else:
            print(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
        if len(missing_keys) > 0:
            print(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path} and are newly initialized: {missing_keys}\nYou should probably"
                " TRAIN this model on a down-stream task to be able to use it for predictions and inference."
            )
        else:
            print(
                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at"
                f" {pretrained_model_name_or_path}.\nIf your task is similar to the task the model of the checkpoint"
                f" was trained on, you can already use {model.__class__.__name__} for predictions without further"
                " training."
            )
        if len(error_msgs) > 0:
            raise RuntimeError(
                "Error(s) in loading state_dict for {}:\n\t{}".format(
                    model.__class__.__name__, "\n\t".join(error_msgs)
                )
            )
        
        model.eval()

        return model

    def forward(
        self,
        images,
        image_shapes,
        gt_boxes=None,
        proposals=None,
        scales_yx=None,
        **kwargs,
    ):
        
        if self.training:
            raise NotImplementedError()
        return self.inference(
            images=images,
            image_shapes=image_shapes,
            gt_boxes=gt_boxes,
            proposals=proposals,
            scales_yx=scales_yx,
            **kwargs,
        )

    @torch.no_grad()
    def inference(
        self,
        images,
        image_shapes,
        gt_boxes=None,
        proposals=None,
        scales_yx=None,
        **kwargs,
    ):
        
        original_sizes = image_shapes * scales_yx
        features = self.backbone(images)

        
        if proposals is None:
            proposal_boxes, _ = self.proposal_generator(images, image_shapes, features, gt_boxes)
        else:
            assert proposals is not None

        
        obj_logits, attr_logits, box_deltas, feature_pooled = self.roi_heads(features, proposal_boxes, gt_boxes)

        
        boxes, classes, class_probs, attrs, attr_probs, roi_features = self.roi_outputs(
            obj_logits=obj_logits,
            attr_logits=attr_logits,
            box_deltas=box_deltas,
            pred_boxes=proposal_boxes,
            features=feature_pooled,
            sizes=image_shapes,
            scales=scales_yx,
        )

        
        subset_kwargs = {
            "max_detections": kwargs.get("max_detections", None),
            "return_tensors": kwargs.get("return_tensors", None),
            "pad_value": kwargs.get("pad_value", 0),
            "padding": kwargs.get("padding", None),
        }
        preds_per_image = torch.tensor([p.size(0) for p in boxes])
        boxes = pad_list_tensors(boxes, preds_per_image, **subset_kwargs)
        classes = pad_list_tensors(classes, preds_per_image, **subset_kwargs)
        class_probs = pad_list_tensors(class_probs, preds_per_image, **subset_kwargs)
        attrs = pad_list_tensors(attrs, preds_per_image, **subset_kwargs)
        attr_probs = pad_list_tensors(attr_probs, preds_per_image, **subset_kwargs)
        roi_features = pad_list_tensors(roi_features, preds_per_image, **subset_kwargs)
        subset_kwargs["padding"] = None
        preds_per_image = pad_list_tensors(preds_per_image, None, **subset_kwargs)
        sizes = pad_list_tensors(image_shapes, None, **subset_kwargs)
        normalized_boxes = norm_box(boxes, original_sizes)
        return OrderedDict(
            {
                "obj_ids": classes,
                "obj_probs": class_probs,
                "attr_ids": attrs,
                "attr_probs": attr_probs,
                "boxes": boxes,
                "sizes": sizes,
                "preds_per_image": preds_per_image,
                "roi_features": roi_features,
                "normalized_boxes": normalized_boxes,
            }
        )