codingutf8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license image processor class for mask2former import math import warnings from typing import any dict iterable list optional set tuple union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import paddingmode getresizeoutputimagesize pad rescale resize tochanneldimensionformat from imageutils import channeldimension imageinput pilimageresampling getimagesize inferchanneldimensionformat isbatched isscaledimage tonumpyarray validimages from utils import imagenetdefaultmean imagenetdefaultstd tensortype istorchavailable istorchtensor logging logger logging getloggername if istorchavailable import torch from torch import nn copied from transformers models detr imageprocessingdetr maxacrossindices def maxacrossindicesvalues iterableany listany return maxvaluesi for valuesi in zipvalues copied from transformers models detr imageprocessingdetr getmaxheightwidth def getmaxheightwidth images listnp ndarray inputdataformat optionalunionstr channeldimension none listint if inputdataformat is none inputdataformat inferchanneldimensionformatimages0 if inputdataformat channeldimension first maxheight maxwidth maxacrossindicesimg shape for img in images elif inputdataformat channeldimension last maxheight maxwidth maxacrossindicesimg shape for img in images else raise valueerrorfinvalid channel dimension format inputdataformat return maxheight maxwidth copied from transformers models detr imageprocessingdetr makepixelmask def makepixelmask image np ndarray outputsize tupleint int inputdataformat optionalunionstr channeldimension none np ndarray inputheight inputwidth getimagesizeimage channeldiminputdataformat mask np zerosoutputsize dtypenp int64 mask inputheight inputwidth 1 return mask copied from transformers models detr imageprocessingdetr binarymasktorle def binarymasktorlemask if istorchtensormask mask mask numpy pixels mask flatten pixels np concatenate0 pixels 0 runs np wherepixels1 pixels 10 1 runs1 2 runs 2 return listruns copied from transformers models detr imageprocessingdetr convertsegmentationtorle def convertsegmentationtorlesegmentation segmentids torch uniquesegmentation runlengthencodings for idx in segmentids mask torch wheresegmentation idx 1 0 rle binarymasktorlemask runlengthencodings appendrle return runlengthencodings copied from transformers models detr imageprocessingdetr removelowandnoobjects def removelowandnoobjectsmasks scores labels objectmaskthreshold numlabels if not masks shape0 scores shape0 labels shape0 raise valueerrormask scores and labels must have the same shape tokeep labels nenumlabels scores objectmaskthreshold return maskstokeep scorestokeep labelstokeep copied from transformers models detr imageprocessingdetr checksegmentvalidity def checksegmentvaliditymasklabels maskprobs k maskthreshold0 5 overlapmaskareathreshold0 8 get the mask associated with the k class maskk masklabels k maskkarea maskk sum compute the area of all the stuff in query k originalarea maskprobsk maskthreshold sum maskexists maskkarea 0 and originalarea 0 eliminate disconnected tiny segments if maskexists arearatio maskkarea originalarea if not arearatio item overlapmaskareathreshold maskexists false return maskexists maskk copied from transformers models detr imageprocessingdetr computesegments def computesegments maskprobs predscores predlabels maskthreshold float 0 5 overlapmaskareathreshold float 0 8 labelidstofuse optionalsetint none targetsize tupleint int none height maskprobs shape1 if targetsize is none else targetsize0 width maskprobs shape2 if targetsize is none else targetsize1 segmentation torch zerosheight width dtypetorch int32 devicemaskprobs device segments listdict if targetsize is not none maskprobs nn functional interpolate maskprobs unsqueeze0 sizetargetsize modebilinear aligncornersfalse 0 currentsegmentid 0 weigh each mask by its prediction score maskprobs predscores view1 1 1 masklabels maskprobs argmax0 height width keep track of instances of each class stuffmemorylist dictstr int for k in rangepredlabels shape0 predclass predlabelsk item shouldfuse predclass in labelidstofuse check if mask exists and large enough to be a segment maskexists maskk checksegmentvalidity masklabels maskprobs k maskthreshold overlapmaskareathreshold if maskexists if predclass in stuffmemorylist currentsegmentid stuffmemorylistpredclass else currentsegmentid 1 add current object segment to final segmentation map segmentationmaskk currentsegmentid segmentscore roundpredscoresk item 6 segments append id currentsegmentid labelid predclass wasfused shouldfuse score segmentscore if shouldfuse stuffmemorylistpredclass currentsegmentid return segmentation segments todo amy move to imagetransforms copied from transformers models maskformer imageprocessingmaskformer convertsegmentationmaptobinarymasks def convertsegmentationmaptobinarymasks segmentationmap np ndarray instanceidtosemanticid optionaldictint int none ignoreindex optionalint none reducelabels bool false if reducelabels and ignoreindex is none raise valueerrorif reducelabels is true ignoreindex must be provided if reducelabels segmentationmap np wheresegmentationmap 0 ignoreindex segmentationmap 1 get unique ids class or instance ids based on input alllabels np uniquesegmentationmap drop background label if applicable if ignoreindex is not none alllabels alllabelsalllabels ignoreindex generate a binary mask for each object instance binarymasks segmentationmap i for i in alllabels binarymasks np stackbinarymasks axis0 numlabels height width convert instance ids to class ids if instanceidtosemanticid is not none labels np zerosalllabels shape0 for label in alllabels classid instanceidtosemanticidlabel 1 if reducelabels else label labelsalllabels label classid 1 if reducelabels else classid else labels alllabels return binarymasks astypenp float32 labels astypenp int64 copied from transformers models maskformer imageprocessingmaskformer getmaskformerresizeoutputimagesize with maskformermask2former def getmask2formerresizeoutputimagesize image np ndarray size unionint tupleint int listint tupleint maxsize optionalint none sizedivisor int 0 defaulttosquare bool true inputdataformat optionalunionstr channeldimension none tupleint int outputsize getresizeoutputimagesize inputimageimage sizesize defaulttosquaredefaulttosquare maxsizemaxsize inputdataformatinputdataformat if sizedivisor 0 height width outputsize height intmath ceilheight sizedivisor sizedivisor width intmath ceilwidth sizedivisor sizedivisor outputsize height width return outputsize class mask2formerimageprocessorbaseimageprocessor r constructs a mask2former image processor the image processor can be used to prepare images and optional targets for the model this image processor inherits from baseimageprocessor which contains most of the main methods users should refer to this superclass for more information regarding those methods args doresize bool optional defaults to true whether to resize the input to a certain size size int optional defaults to 800 resize the input to the given size only has an effect if doresize is set to true if size is a sequence like width height output size will be matched to this if size is an int smaller edge of the image will be matched to this number i e if height width then image will be rescaled to size height width size sizedivisor int optional defaults to 32 some backbones need images divisible by a certain number if not passed it defaults to the value used in swin transformer resample int optional defaults to resampling bilinear an optional resampling filter this can be one of pil image resampling nearest pil image resampling box pil image resampling bilinear pil image resampling hamming pil image resampling bicubic or pil image resampling lanczos only has an effect if doresize is set to true dorescale bool optional defaults to true whether to rescale the input to a certain scale rescalefactor float optional defaults to 1 255 rescale the input by the given factor only has an effect if dorescale is set to true donormalize bool optional defaults to true whether or not to normalize the input with mean and standard deviation imagemean int optional defaults to 0 485 0 456 0 406 the sequence of means for each channel to be used when normalizing images defaults to the imagenet mean imagestd int optional defaults to 0 229 0 224 0 225 the sequence of standard deviations for each channel to be used when normalizing images defaults to the imagenet std ignoreindex int optional label to be assigned to background pixels in segmentation maps if provided segmentation map pixels denoted with 0 background will be replaced with ignoreindex reducelabels bool optional defaults to false whether or not to decrement all label values of segmentation maps by 1 usually used for datasets where 0 is used for background and background itself is not included in all classes of a dataset e g ade20k the background label will be replaced by ignoreindex we make maxsize a private attribute so we can pass it as a default value in the preprocess method whilst size can still be pass in as an int overrides the fromdict method from the base class to make sure parameters are updated if image processor is created using fromdict and kwargs e g mask2formerimageprocessor frompretrainedcheckpoint maxsize800 copied from transformers models maskformer imageprocessingmaskformer maskformerimageprocessor resize with getmaskformerresizeoutputimagesizegetmask2formerresizeoutputimagesize resize the image to the given size size can be minsize scalar or height width tuple if size is an int smaller edge of the image will be matched to this number args image np ndarray image to resize size dictstr int the size of the output image sizedivisor int optional defaults to 0 if sizedivisor is given the output image size will be divisible by the number resample pilimageresampling resampling filter optional defaults to pilimageresampling bilinear resampling filter to use when resizing the image dataformat channeldimension or str optional the channel dimension format for the output image if unset the channel dimension format of the input image is used inputdataformat channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred copied from transformers models detr imageprocessingdetr detrimageprocessor rescale rescale the image by the given factor image image rescalefactor args image np ndarray image to rescale rescalefactor float the value to use for rescaling dataformat str or channeldimension optional the channel dimension format for the output image if unset the channel dimension format of the input image is used can be one of channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format inputdataformat str or channeldimension optional the channel dimension format for the input image if unset is inferred from the input image can be one of channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format copied from transformers models maskformer imageprocessingmaskformer maskformerimageprocessor convertsegmentationmaptobinarymasks preprocesses a single image all transformations expect numpy arrays image tonumpyarrayimage if isscaledimageimage and dorescale logger warningonce it looks like you are trying to rescale already rescaled images if the input images have pixel values between 0 and 1 set dorescalefalse to avoid rescaling them again if inputdataformat is none inputdataformat inferchanneldimensionformatimage image self preprocess imageimage doresizedoresize sizesize sizedivisorsizedivisor resampleresample dorescaledorescale rescalefactorrescalefactor donormalizedonormalize imagemeanimagemean imagestdimagestd inputdataformatinputdataformat if dataformat is not none image tochanneldimensionformatimage dataformat inputchanneldiminputdataformat return image def preprocessmask self segmentationmap imageinput doresize bool none size dictstr int none sizedivisor int 0 inputdataformat optionalunionstr channeldimension none np ndarray add channel dimension if missing needed for certain transformations todo amy remork segmentation map processing to include reducing labels and resizing which doesn t drop segment ids 255 remove extra channel dimension if added for processing copied from transformers models detr imageprocessingdetr detrimageprocessor padimage pad an image with zeros to the given size copied from transformers models detr imageprocessingdetr detrimageprocessor pad pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width in the batch and optionally returns their corresponding pixel mask args image np ndarray image to pad constantvalues float or iterablefloat optional the value to use for the padding if mode is constant returnpixelmask bool optional defaults to true whether to return a pixel mask returntensors str or tensortype optional the type of tensors to return can be one of unset return a list of np ndarray tensortype tensorflow or tf return a batch of type tf tensor tensortype pytorch or pt return a batch of type torch tensor tensortype numpy or np return a batch of type np ndarray tensortype jax or jax return a batch of type jax numpy ndarray dataformat str or channeldimension optional the channel dimension format of the image if not provided it will be the same as the input image inputdataformat channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred pad images up to the largest image in a batch and create a corresponding pixelmask mask2former addresses semantic segmentation with a mask classification paradigm thus input segmentation maps will be converted to lists of binary masks and their respective labels let s see an example assuming segmentationmaps 2 6 7 9 the output will contain masklabels 1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1 four binary masks and classlabels 2 6 7 9 the labels for each mask args pixelvalueslist listimageinput list of images pixel values to be padded each image should be a tensor of shape channels height width segmentationmaps imageinput optional the corresponding semantic segmentation maps with the pixelwise annotations bool optional defaults to true whether or not to pad images up to the largest image in a batch and create a pixel mask if left to the default will return a pixel mask that is 1 for pixels that are real i e not masked 0 for pixels that are padding i e masked instanceidtosemanticid listdictint int or dictint int optional a mapping between object instance ids and class ids if passed segmentationmaps is treated as an instance segmentation map where each pixel represents an instance id can be provided as a single dictionary with a globaldatasetlevel mapping or as a list of dictionaries one per image to map instance ids in each image separately returntensors str or fileutils tensortype optional if set will return tensors instead of numpy arrays if set to pt return pytorch torch tensor objects inputdataformat channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred returns batchfeature a batchfeature with the following fields pixelvalues pixel values to be fed to a model pixelmask pixel mask to be fed to a model when true or if pixelmask is in self modelinputnames masklabels optional list of mask labels of shape labels height width to be fed to a model when annotations are provided classlabels optional list of class labels of shape labels to be fed to a model when annotations are provided they identify the labels of masklabels e g the label of masklabelsij if classlabelsij convert to list of binary masks and labels use instance2classid mapping per image we add an axis to make them compatible with the transformations library this will be removed in the future we cannot batch them since they don t share a common class size converts the output of mask2formerforuniversalsegmentation into semantic segmentation maps only supports pytorch args outputs mask2formerforuniversalsegmentation raw outputs of the model targetsizes listtupleint int optional list of length batchsize where each list item tupleint int corresponds to the requested final size height width of each prediction if left to none predictions will not be resized returns listtorch tensor a list of length batchsize where each item is a semantic segmentation map of shape height width corresponding to the targetsizes entry if targetsizes is specified each entry of each torch tensor correspond to a semantic class id scale back to preprocessed image size 384 384 for all models remove the null class 1 semantic segmentation logits of shape batchsize numclasses height width resize logits and compute semantic segmentation maps converts the output of mask2formerforuniversalsegmentationoutput into instance segmentation predictions only supports pytorch args outputs mask2formerforuniversalsegmentation raw outputs of the model threshold float optional defaults to 0 5 the probability score threshold to keep predicted instance masks maskthreshold float optional defaults to 0 5 threshold to use when turning the predicted masks into binary values overlapmaskareathreshold float optional defaults to 0 8 the overlap mask area threshold to merge or discard small disconnected parts within each binary instance mask targetsizes listtuple optional list of length batchsize where each list item tupleint int corresponds to the requested final size height width of each prediction if left to none predictions will not be resized returncocoannotation bool optional defaults to false if set to true segmentation maps are returned in coco runlength encoding rle format returnbinarymaps bool optional defaults to false if set to true segmentation maps are returned as a concatenated tensor of binary segmentation maps one per detected instance returns listdict a list of dictionaries one per image each dictionary containing two keys segmentation a tensor of shape height width where each pixel represents a segmentid or listlist runlength encoding rle of the segmentation map if returncocoannotation is set to true set to none if no mask if found above threshold segmentsinfo a dictionary that contains additional information on each segment id an integer representing the segmentid labelid an integer representing the label semantic class id corresponding to segmentid score prediction score of segment with segmentid batchsize numqueries numclasses1 batchsize numqueries height width scale back to preprocessed image size 384 384 for all models loop over items in batch size calculate average mask prob return segmentation map in runlength encoding rle format return a concatenated tensor of binary instance maps converts the output of mask2formerforuniversalsegmentationoutput into image panoptic segmentation predictions only supports pytorch args outputs mask2formerforuniversalsegmentationoutput the outputs from mask2formerforuniversalsegmentation threshold float optional defaults to 0 5 the probability score threshold to keep predicted instance masks maskthreshold float optional defaults to 0 5 threshold to use when turning the predicted masks into binary values overlapmaskareathreshold float optional defaults to 0 8 the overlap mask area threshold to merge or discard small disconnected parts within each binary instance mask labelidstofuse setint optional the labels in this state will have all their instances be fused together for instance we could say there can only be one sky in an image but several persons so the label id for sky would be in that set but not the one for person targetsizes listtuple optional list of length batchsize where each list item tupleint int corresponds to the requested final size height width of each prediction in batch if left to none predictions will not be resized returns listdict a list of dictionaries one per image each dictionary containing two keys segmentation a tensor of shape height width where each pixel represents a segmentid set to none if no mask if found above threshold if targetsizes is specified segmentation is resized to the corresponding targetsizes entry segmentsinfo a dictionary that contains additional information on each segment id an integer representing the segmentid labelid an integer representing the label semantic class id corresponding to segmentid wasfused a boolean true if labelid was in labelidstofuse false otherwise multiple instances of the same class label were fused and assigned a single segmentid score prediction score of segment with segmentid scale back to preprocessed image size 384 384 for all models predicted label and score of each query batchsize numqueries loop over items in batch size no mask found get segmentation map and segment information of batch item coding utf 8 2022 the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license image processor class for mask2former copied from transformers models detr image_processing_detr max_across_indices return the maximum value across all indices of an iterable of values copied from transformers models detr image_processing_detr get_max_height_width get the maximum height and width across all images in a batch copied from transformers models detr image_processing_detr make_pixel_mask make a pixel mask for the image where 1 indicates a valid pixel and 0 indicates padding args image np ndarray image to make the pixel mask for output_size tuple int int output size of the mask copied from transformers models detr image_processing_detr binary_mask_to_rle converts given binary mask of shape height width to the run length encoding rle format args mask torch tensor or numpy array a binary mask tensor of shape height width where 0 denotes background and 1 denotes the target segment_id or class_id returns list run length encoded list of the binary mask refer to coco api for more information about the rle format copied from transformers models detr image_processing_detr convert_segmentation_to_rle converts given segmentation map of shape height width to the run length encoding rle format args segmentation torch tensor or numpy array a segmentation map of shape height width where each value denotes a segment or class id returns list list a list of lists where each list is the run length encoding of a segment class id copied from transformers models detr image_processing_detr remove_low_and_no_objects binarize the given masks using object_mask_threshold it returns the associated values of masks scores and labels args masks torch tensor a tensor of shape num_queries height width scores torch tensor a tensor of shape num_queries labels torch tensor a tensor of shape num_queries object_mask_threshold float a number between 0 and 1 used to binarize the masks raises valueerror raised when the first dimension doesn t match in all input tensors returns tuple torch tensor torch tensor torch tensor the masks scores and labels without the region object_mask_threshold copied from transformers models detr image_processing_detr check_segment_validity get the mask associated with the k class compute the area of all the stuff in query k eliminate disconnected tiny segments copied from transformers models detr image_processing_detr compute_segments weigh each mask by its prediction score height width keep track of instances of each class check if mask exists and large enough to be a segment add current object segment to final segmentation map todo amy move to image_transforms copied from transformers models maskformer image_processing_maskformer convert_segmentation_map_to_binary_masks get unique ids class or instance ids based on input drop background label if applicable generate a binary mask for each object instance num_labels height width convert instance ids to class ids copied from transformers models maskformer image_processing_maskformer get_maskformer_resize_output_image_size with maskformer mask2former computes the output size given the desired size args image np ndarray the input image size int or tuple int int or list int or tuple int the size of the output image max_size int optional the maximum size of the output image size_divisor int optional defaults to 0 if size_divisor is given the output image size will be divisible by the number default_to_square bool optional defaults to true whether to default to square if no size is provided input_data_format channeldimension or str optional the channel dimension format of the input image if unset will use the inferred format from the input returns tuple int int the output size constructs a mask2former image processor the image processor can be used to prepare image s and optional targets for the model this image processor inherits from baseimageprocessor which contains most of the main methods users should refer to this superclass for more information regarding those methods args do_resize bool optional defaults to true whether to resize the input to a certain size size int optional defaults to 800 resize the input to the given size only has an effect if do_resize is set to true if size is a sequence like width height output size will be matched to this if size is an int smaller edge of the image will be matched to this number i e if height width then image will be rescaled to size height width size size_divisor int optional defaults to 32 some backbones need images divisible by a certain number if not passed it defaults to the value used in swin transformer resample int optional defaults to resampling bilinear an optional resampling filter this can be one of pil image resampling nearest pil image resampling box pil image resampling bilinear pil image resampling hamming pil image resampling bicubic or pil image resampling lanczos only has an effect if do_resize is set to true do_rescale bool optional defaults to true whether to rescale the input to a certain scale rescale_factor float optional defaults to 1 255 rescale the input by the given factor only has an effect if do_rescale is set to true do_normalize bool optional defaults to true whether or not to normalize the input with mean and standard deviation image_mean int optional defaults to 0 485 0 456 0 406 the sequence of means for each channel to be used when normalizing images defaults to the imagenet mean image_std int optional defaults to 0 229 0 224 0 225 the sequence of standard deviations for each channel to be used when normalizing images defaults to the imagenet std ignore_index int optional label to be assigned to background pixels in segmentation maps if provided segmentation map pixels denoted with 0 background will be replaced with ignore_index reduce_labels bool optional defaults to false whether or not to decrement all label values of segmentation maps by 1 usually used for datasets where 0 is used for background and background itself is not included in all classes of a dataset e g ade20k the background label will be replaced by ignore_index we make max_size a private attribute so we can pass it as a default value in the preprocess method whilst size can still be pass in as an int overrides the from_dict method from the base class to make sure parameters are updated if image processor is created using from_dict and kwargs e g mask2formerimageprocessor from_pretrained checkpoint max_size 800 copied from transformers models maskformer image_processing_maskformer maskformerimageprocessor resize with get_maskformer_resize_output_image_size get_mask2former_resize_output_image_size resize the image to the given size size can be min_size scalar or height width tuple if size is an int smaller edge of the image will be matched to this number args image np ndarray image to resize size dict str int the size of the output image size_divisor int optional defaults to 0 if size_divisor is given the output image size will be divisible by the number resample pilimageresampling resampling filter optional defaults to pilimageresampling bilinear resampling filter to use when resizing the image data_format channeldimension or str optional the channel dimension format for the output image if unset the channel dimension format of the input image is used input_data_format channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred copied from transformers models detr image_processing_detr detrimageprocessor rescale rescale the image by the given factor image image rescale_factor args image np ndarray image to rescale rescale_factor float the value to use for rescaling data_format str or channeldimension optional the channel dimension format for the output image if unset the channel dimension format of the input image is used can be one of channels_first or channeldimension first image in num_channels height width format channels_last or channeldimension last image in height width num_channels format input_data_format str or channeldimension optional the channel dimension format for the input image if unset is inferred from the input image can be one of channels_first or channeldimension first image in num_channels height width format channels_last or channeldimension last image in height width num_channels format copied from transformers models maskformer image_processing_maskformer maskformerimageprocessor convert_segmentation_map_to_binary_masks preprocesses a single image all transformations expect numpy arrays preprocesses a single mask add channel dimension if missing needed for certain transformations todo amy remork segmentation map processing to include reducing labels and resizing which doesn t drop segment ids 255 remove extra channel dimension if added for processing copied from transformers models detr image_processing_detr detrimageprocessor _pad_image pad an image with zeros to the given size copied from transformers models detr image_processing_detr detrimageprocessor pad pads a batch of images to the bottom and right of the image with zeros to the size of largest height and width in the batch and optionally returns their corresponding pixel mask args image np ndarray image to pad constant_values float or iterable float optional the value to use for the padding if mode is constant return_pixel_mask bool optional defaults to true whether to return a pixel mask return_tensors str or tensortype optional the type of tensors to return can be one of unset return a list of np ndarray tensortype tensorflow or tf return a batch of type tf tensor tensortype pytorch or pt return a batch of type torch tensor tensortype numpy or np return a batch of type np ndarray tensortype jax or jax return a batch of type jax numpy ndarray data_format str or channeldimension optional the channel dimension format of the image if not provided it will be the same as the input image input_data_format channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred pad images up to the largest image in a batch and create a corresponding pixel_mask mask2former addresses semantic segmentation with a mask classification paradigm thus input segmentation maps will be converted to lists of binary masks and their respective labels let s see an example assuming segmentation_maps 2 6 7 9 the output will contain mask_labels 1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1 four binary masks and class_labels 2 6 7 9 the labels for each mask args pixel_values_list list imageinput list of images pixel values to be padded each image should be a tensor of shape channels height width segmentation_maps imageinput optional the corresponding semantic segmentation maps with the pixel wise annotations bool optional defaults to true whether or not to pad images up to the largest image in a batch and create a pixel mask if left to the default will return a pixel mask that is 1 for pixels that are real i e not masked 0 for pixels that are padding i e masked instance_id_to_semantic_id list dict int int or dict int int optional a mapping between object instance ids and class ids if passed segmentation_maps is treated as an instance segmentation map where each pixel represents an instance id can be provided as a single dictionary with a global dataset level mapping or as a list of dictionaries one per image to map instance ids in each image separately return_tensors str or file_utils tensortype optional if set will return tensors instead of numpy arrays if set to pt return pytorch torch tensor objects input_data_format channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred returns batchfeature a batchfeature with the following fields pixel_values pixel values to be fed to a model pixel_mask pixel mask to be fed to a model when true or if pixel_mask is in self model_input_names mask_labels optional list of mask labels of shape labels height width to be fed to a model when annotations are provided class_labels optional list of class labels of shape labels to be fed to a model when annotations are provided they identify the labels of mask_labels e g the label of mask_labels i j if class_labels i j convert to list of binary masks and labels use instance2class_id mapping per image we add an axis to make them compatible with the transformations library this will be removed in the future we cannot batch them since they don t share a common class size converts the output of mask2formerforuniversalsegmentation into semantic segmentation maps only supports pytorch args outputs mask2formerforuniversalsegmentation raw outputs of the model target_sizes list tuple int int optional list of length batch_size where each list item tuple int int corresponds to the requested final size height width of each prediction if left to none predictions will not be resized returns list torch tensor a list of length batch_size where each item is a semantic segmentation map of shape height width corresponding to the target_sizes entry if target_sizes is specified each entry of each torch tensor correspond to a semantic class id batch_size num_queries num_classes 1 batch_size num_queries height width scale back to preprocessed image size 384 384 for all models remove the null class 1 batch_size num_queries height width semantic segmentation logits of shape batch_size num_classes height width resize logits and compute semantic segmentation maps converts the output of mask2formerforuniversalsegmentationoutput into instance segmentation predictions only supports pytorch args outputs mask2formerforuniversalsegmentation raw outputs of the model threshold float optional defaults to 0 5 the probability score threshold to keep predicted instance masks mask_threshold float optional defaults to 0 5 threshold to use when turning the predicted masks into binary values overlap_mask_area_threshold float optional defaults to 0 8 the overlap mask area threshold to merge or discard small disconnected parts within each binary instance mask target_sizes list tuple optional list of length batch_size where each list item tuple int int corresponds to the requested final size height width of each prediction if left to none predictions will not be resized return_coco_annotation bool optional defaults to false if set to true segmentation maps are returned in coco run length encoding rle format return_binary_maps bool optional defaults to false if set to true segmentation maps are returned as a concatenated tensor of binary segmentation maps one per detected instance returns list dict a list of dictionaries one per image each dictionary containing two keys segmentation a tensor of shape height width where each pixel represents a segment_id or list list run length encoding rle of the segmentation map if return_coco_annotation is set to true set to none if no mask if found above threshold segments_info a dictionary that contains additional information on each segment id an integer representing the segment_id label_id an integer representing the label semantic class id corresponding to segment_id score prediction score of segment with segment_id batch_size num_queries num_classes 1 batch_size num_queries height width scale back to preprocessed image size 384 384 for all models loop over items in batch size calculate average mask prob return segmentation map in run length encoding rle format return a concatenated tensor of binary instance maps converts the output of mask2formerforuniversalsegmentationoutput into image panoptic segmentation predictions only supports pytorch args outputs mask2formerforuniversalsegmentationoutput the outputs from mask2formerforuniversalsegmentation threshold float optional defaults to 0 5 the probability score threshold to keep predicted instance masks mask_threshold float optional defaults to 0 5 threshold to use when turning the predicted masks into binary values overlap_mask_area_threshold float optional defaults to 0 8 the overlap mask area threshold to merge or discard small disconnected parts within each binary instance mask label_ids_to_fuse set int optional the labels in this state will have all their instances be fused together for instance we could say there can only be one sky in an image but several persons so the label id for sky would be in that set but not the one for person target_sizes list tuple optional list of length batch_size where each list item tuple int int corresponds to the requested final size height width of each prediction in batch if left to none predictions will not be resized returns list dict a list of dictionaries one per image each dictionary containing two keys segmentation a tensor of shape height width where each pixel represents a segment_id set to none if no mask if found above threshold if target_sizes is specified segmentation is resized to the corresponding target_sizes entry segments_info a dictionary that contains additional information on each segment id an integer representing the segment_id label_id an integer representing the label semantic class id corresponding to segment_id was_fused a boolean true if label_id was in label_ids_to_fuse false otherwise multiple instances of the same class label were fused and assigned a single segment_id score prediction score of segment with segment_id batch_size num_queries num_classes 1 batch_size num_queries height width scale back to preprocessed image size 384 384 for all models batch_size num_queries height width predicted label and score of each query batch_size num_queries loop over items in batch size no mask found get segmentation map and segment information of batch item
import math
import warnings
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    PaddingMode,
    get_resize_output_image_size,
    pad,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    infer_channel_dimension_format,
    is_batched,
    is_scaled_image,
    to_numpy_array,
    valid_images,
)
from ...utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    TensorType,
    is_torch_available,
    is_torch_tensor,
    logging,
)


logger = logging.get_logger(__name__)


if is_torch_available():
    import torch
    from torch import nn



def max_across_indices(values: Iterable[Any]) -> List[Any]:
    
    return [max(values_i) for values_i in zip(*values)]



def get_max_height_width(
    images: List[np.ndarray], input_data_format: Optional[Union[str, ChannelDimension]] = None
) -> List[int]:
    
    if input_data_format is None:
        input_data_format = infer_channel_dimension_format(images[0])

    if input_data_format == ChannelDimension.FIRST:
        _, max_height, max_width = max_across_indices([img.shape for img in images])
    elif input_data_format == ChannelDimension.LAST:
        max_height, max_width, _ = max_across_indices([img.shape for img in images])
    else:
        raise ValueError(f"Invalid channel dimension format: {input_data_format}")
    return (max_height, max_width)



def make_pixel_mask(
    image: np.ndarray, output_size: Tuple[int, int], input_data_format: Optional[Union[str, ChannelDimension]] = None
) -> np.ndarray:
    
    input_height, input_width = get_image_size(image, channel_dim=input_data_format)
    mask = np.zeros(output_size, dtype=np.int64)
    mask[:input_height, :input_width] = 1
    return mask



def binary_mask_to_rle(mask):
    
    if is_torch_tensor(mask):
        mask = mask.numpy()

    pixels = mask.flatten()
    pixels = np.concatenate([[0], pixels, [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    return list(runs)



def convert_segmentation_to_rle(segmentation):
    
    segment_ids = torch.unique(segmentation)

    run_length_encodings = []
    for idx in segment_ids:
        mask = torch.where(segmentation == idx, 1, 0)
        rle = binary_mask_to_rle(mask)
        run_length_encodings.append(rle)

    return run_length_encodings



def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
    
    if not (masks.shape[0] == scores.shape[0] == labels.shape[0]):
        raise ValueError("mask, scores and labels must have the same shape!")

    to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)

    return masks[to_keep], scores[to_keep], labels[to_keep]



def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8):
    
    mask_k = mask_labels == k
    mask_k_area = mask_k.sum()

    
    original_area = (mask_probs[k] >= mask_threshold).sum()
    mask_exists = mask_k_area > 0 and original_area > 0

    
    if mask_exists:
        area_ratio = mask_k_area / original_area
        if not area_ratio.item() > overlap_mask_area_threshold:
            mask_exists = False

    return mask_exists, mask_k



def compute_segments(
    mask_probs,
    pred_scores,
    pred_labels,
    mask_threshold: float = 0.5,
    overlap_mask_area_threshold: float = 0.8,
    label_ids_to_fuse: Optional[Set[int]] = None,
    target_size: Tuple[int, int] = None,
):
    height = mask_probs.shape[1] if target_size is None else target_size[0]
    width = mask_probs.shape[2] if target_size is None else target_size[1]

    segmentation = torch.zeros((height, width), dtype=torch.int32, device=mask_probs.device)
    segments: List[Dict] = []

    if target_size is not None:
        mask_probs = nn.functional.interpolate(
            mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False
        )[0]

    current_segment_id = 0

    
    mask_probs *= pred_scores.view(-1, 1, 1)
    mask_labels = mask_probs.argmax(0)  

    
    stuff_memory_list: Dict[str, int] = {}
    for k in range(pred_labels.shape[0]):
        pred_class = pred_labels[k].item()
        should_fuse = pred_class in label_ids_to_fuse

        
        mask_exists, mask_k = check_segment_validity(
            mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
        )

        if mask_exists:
            if pred_class in stuff_memory_list:
                current_segment_id = stuff_memory_list[pred_class]
            else:
                current_segment_id += 1

            
            segmentation[mask_k] = current_segment_id
            segment_score = round(pred_scores[k].item(), 6)
            segments.append(
                {
                    "id": current_segment_id,
                    "label_id": pred_class,
                    "was_fused": should_fuse,
                    "score": segment_score,
                }
            )
            if should_fuse:
                stuff_memory_list[pred_class] = current_segment_id

    return segmentation, segments




def convert_segmentation_map_to_binary_masks(
    segmentation_map: "np.ndarray",
    instance_id_to_semantic_id: Optional[Dict[int, int]] = None,
    ignore_index: Optional[int] = None,
    reduce_labels: bool = False,
):
    if reduce_labels and ignore_index is None:
        raise ValueError("If `reduce_labels` is True, `ignore_index` must be provided.")

    if reduce_labels:
        segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1)

    
    all_labels = np.unique(segmentation_map)

    
    if ignore_index is not None:
        all_labels = all_labels[all_labels != ignore_index]

    
    binary_masks = [(segmentation_map == i) for i in all_labels]
    binary_masks = np.stack(binary_masks, axis=0)  

    
    if instance_id_to_semantic_id is not None:
        labels = np.zeros(all_labels.shape[0])

        for label in all_labels:
            class_id = instance_id_to_semantic_id[label + 1 if reduce_labels else label]
            labels[all_labels == label] = class_id - 1 if reduce_labels else class_id
    else:
        labels = all_labels

    return binary_masks.astype(np.float32), labels.astype(np.int64)



def get_mask2former_resize_output_image_size(
    image: np.ndarray,
    size: Union[int, Tuple[int, int], List[int], Tuple[int]],
    max_size: Optional[int] = None,
    size_divisor: int = 0,
    default_to_square: bool = True,
    input_data_format: Optional[Union[str, ChannelDimension]] = None,
) -> Tuple[int, int]:
    
    output_size = get_resize_output_image_size(
        input_image=image,
        size=size,
        default_to_square=default_to_square,
        max_size=max_size,
        input_data_format=input_data_format,
    )

    if size_divisor > 0:
        height, width = output_size
        height = int(math.ceil(height / size_divisor) * size_divisor)
        width = int(math.ceil(width / size_divisor) * size_divisor)
        output_size = (height, width)

    return output_size


class Mask2FormerImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values", "pixel_mask"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        size_divisor: int = 32,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        do_rescale: bool = True,
        rescale_factor: float = 1 / 255,
        do_normalize: bool = True,
        image_mean: Union[float, List[float]] = None,
        image_std: Union[float, List[float]] = None,
        ignore_index: Optional[int] = None,
        reduce_labels: bool = False,
        **kwargs,
    ):
        if "size_divisibility" in kwargs:
            warnings.warn(
                "The `size_divisibility` argument is deprecated and will be removed in v4.27. Please use "
                "`size_divisor` instead.",
                FutureWarning,
            )
            size_divisor = kwargs.pop("size_divisibility")
        if "max_size" in kwargs:
            warnings.warn(
                "The `max_size` argument is deprecated and will be removed in v4.27. Please use size['longest_edge']"
                " instead.",
                FutureWarning,
            )
            
            
            self._max_size = kwargs.pop("max_size")
        else:
            self._max_size = 1333

        size = size if size is not None else {"shortest_edge": 800, "longest_edge": self._max_size}
        size = get_size_dict(size, max_size=self._max_size, default_to_square=False)

        super().__init__(**kwargs)
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.size_divisor = size_divisor
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_normalize = do_normalize
        self.image_mean = image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
        self.image_std = image_std if image_std is not None else IMAGENET_DEFAULT_STD
        self.ignore_index = ignore_index
        self.reduce_labels = reduce_labels

    @classmethod
    def from_dict(cls, image_processor_dict: Dict[str, Any], **kwargs):
        
        image_processor_dict = image_processor_dict.copy()
        if "max_size" in kwargs:
            image_processor_dict["max_size"] = kwargs.pop("max_size")
        if "size_divisibility" in kwargs:
            image_processor_dict["size_divisibility"] = kwargs.pop("size_divisibility")
        return super().from_dict(image_processor_dict, **kwargs)

    
    def resize(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        size_divisor: int = 0,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        data_format=None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> np.ndarray:
        
        if "max_size" in kwargs:
            warnings.warn(
                "The `max_size` parameter is deprecated and will be removed in v4.27. "
                "Please specify in `size['longest_edge'] instead`.",
                FutureWarning,
            )
            max_size = kwargs.pop("max_size")
        else:
            max_size = None
        size = get_size_dict(size, max_size=max_size, default_to_square=False)
        if "shortest_edge" in size and "longest_edge" in size:
            size, max_size = size["shortest_edge"], size["longest_edge"]
        elif "height" in size and "width" in size:
            size = (size["height"], size["width"])
            max_size = None
        else:
            raise ValueError(
                "Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got"
                f" {size.keys()}."
            )
        size = get_mask2former_resize_output_image_size(
            image=image,
            size=size,
            max_size=max_size,
            size_divisor=size_divisor,
            default_to_square=False,
            input_data_format=input_data_format,
        )
        image = resize(
            image, size=size, resample=resample, data_format=data_format, input_data_format=input_data_format, **kwargs
        )
        return image

    
    def rescale(
        self,
        image: np.ndarray,
        rescale_factor: float,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        return rescale(image, rescale_factor, data_format=data_format, input_data_format=input_data_format)

    
    def convert_segmentation_map_to_binary_masks(
        self,
        segmentation_map: "np.ndarray",
        instance_id_to_semantic_id: Optional[Dict[int, int]] = None,
        ignore_index: Optional[int] = None,
        reduce_labels: bool = False,
    ):
        reduce_labels = reduce_labels if reduce_labels is not None else self.reduce_labels
        ignore_index = ignore_index if ignore_index is not None else self.ignore_index
        return convert_segmentation_map_to_binary_masks(
            segmentation_map=segmentation_map,
            instance_id_to_semantic_id=instance_id_to_semantic_id,
            ignore_index=ignore_index,
            reduce_labels=reduce_labels,
        )

    def __call__(self, images, segmentation_maps=None, **kwargs) -> BatchFeature:
        return self.preprocess(images, segmentation_maps=segmentation_maps, **kwargs)

    def _preprocess(
        self,
        image: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        size_divisor: int = None,
        resample: PILImageResampling = None,
        do_rescale: bool = None,
        rescale_factor: float = None,
        do_normalize: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        if do_resize:
            image = self.resize(
                image, size=size, size_divisor=size_divisor, resample=resample, input_data_format=input_data_format
            )
        if do_rescale:
            image = self.rescale(image, rescale_factor=rescale_factor, input_data_format=input_data_format)
        if do_normalize:
            image = self.normalize(image, mean=image_mean, std=image_std, input_data_format=input_data_format)
        return image

    def _preprocess_image(
        self,
        image: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        size_divisor: int = None,
        resample: PILImageResampling = None,
        do_rescale: bool = None,
        rescale_factor: float = None,
        do_normalize: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        
        image = to_numpy_array(image)
        if is_scaled_image(image) and do_rescale:
            logger.warning_once(
                "It looks like you are trying to rescale already rescaled images. If the input"
                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
            )
        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(image)
        image = self._preprocess(
            image=image,
            do_resize=do_resize,
            size=size,
            size_divisor=size_divisor,
            resample=resample,
            do_rescale=do_rescale,
            rescale_factor=rescale_factor,
            do_normalize=do_normalize,
            image_mean=image_mean,
            image_std=image_std,
            input_data_format=input_data_format,
        )
        if data_format is not None:
            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
        return image

    def _preprocess_mask(
        self,
        segmentation_map: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        size_divisor: int = 0,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        segmentation_map = to_numpy_array(segmentation_map)
        
        if segmentation_map.ndim == 2:
            added_channel_dim = True
            segmentation_map = segmentation_map[None, ...]
            input_data_format = ChannelDimension.FIRST
        else:
            added_channel_dim = False
            if input_data_format is None:
                input_data_format = infer_channel_dimension_format(segmentation_map)
        
        
        
        segmentation_map = self._preprocess(
            image=segmentation_map,
            do_resize=do_resize,
            resample=PILImageResampling.NEAREST,
            size=size,
            size_divisor=size_divisor,
            do_rescale=False,
            do_normalize=False,
            input_data_format=input_data_format,
        )
        
        if added_channel_dim:
            segmentation_map = segmentation_map.squeeze(0)
        return segmentation_map

    def preprocess(
        self,
        images: ImageInput,
        segmentation_maps: Optional[ImageInput] = None,
        instance_id_to_semantic_id: Optional[Dict[int, int]] = None,
        do_resize: Optional[bool] = None,
        size: Optional[Dict[str, int]] = None,
        size_divisor: Optional[int] = None,
        resample: PILImageResampling = None,
        do_rescale: Optional[bool] = None,
        rescale_factor: Optional[float] = None,
        do_normalize: Optional[bool] = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        ignore_index: Optional[int] = None,
        reduce_labels: Optional[bool] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: Union[str, ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> BatchFeature:
        if "pad_and_return_pixel_mask" in kwargs:
            warnings.warn(
                "The `pad_and_return_pixel_mask` argument is deprecated and will be removed in a future version",
                FutureWarning,
            )

        do_resize = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=False, max_size=self._max_size)
        size_divisor = size_divisor if size_divisor is not None else self.size_divisor
        resample = resample if resample is not None else self.resample
        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        image_mean = image_mean if image_mean is not None else self.image_mean
        image_std = image_std if image_std is not None else self.image_std
        ignore_index = ignore_index if ignore_index is not None else self.ignore_index
        reduce_labels = reduce_labels if reduce_labels is not None else self.reduce_labels

        if do_resize is not None and size is None or size_divisor is None:
            raise ValueError("If `do_resize` is True, `size` and `size_divisor` must be provided.")

        if do_rescale is not None and rescale_factor is None:
            raise ValueError("If `do_rescale` is True, `rescale_factor` must be provided.")

        if do_normalize is not None and (image_mean is None or image_std is None):
            raise ValueError("If `do_normalize` is True, `image_mean` and `image_std` must be provided.")

        if not valid_images(images):
            raise ValueError(
                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
                "torch.Tensor, tf.Tensor or jax.ndarray."
            )

        if segmentation_maps is not None and not valid_images(segmentation_maps):
            raise ValueError(
                "Invalid segmentation map type. Must be of type PIL.Image.Image, numpy.ndarray, "
                "torch.Tensor, tf.Tensor or jax.ndarray."
            )

        if not is_batched(images):
            images = [images]
            segmentation_maps = [segmentation_maps] if segmentation_maps is not None else None

        if segmentation_maps is not None and len(images) != len(segmentation_maps):
            raise ValueError("Images and segmentation maps must have the same length.")

        images = [
            self._preprocess_image(
                image,
                do_resize=do_resize,
                size=size,
                size_divisor=size_divisor,
                resample=resample,
                do_rescale=do_rescale,
                rescale_factor=rescale_factor,
                do_normalize=do_normalize,
                image_mean=image_mean,
                image_std=image_std,
                data_format=data_format,
                input_data_format=input_data_format,
            )
            for image in images
        ]

        if segmentation_maps is not None:
            segmentation_maps = [
                self._preprocess_mask(
                    segmentation_map, do_resize, size, size_divisor, input_data_format=input_data_format
                )
                for segmentation_map in segmentation_maps
            ]
        encoded_inputs = self.encode_inputs(
            images,
            segmentation_maps,
            instance_id_to_semantic_id,
            ignore_index,
            reduce_labels,
            return_tensors,
            input_data_format=input_data_format,
        )
        return encoded_inputs

    
    def _pad_image(
        self,
        image: np.ndarray,
        output_size: Tuple[int, int],
        constant_values: Union[float, Iterable[float]] = 0,
        data_format: Optional[ChannelDimension] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        input_height, input_width = get_image_size(image, channel_dim=input_data_format)
        output_height, output_width = output_size

        pad_bottom = output_height - input_height
        pad_right = output_width - input_width
        padding = ((0, pad_bottom), (0, pad_right))
        padded_image = pad(
            image,
            padding,
            mode=PaddingMode.CONSTANT,
            constant_values=constant_values,
            data_format=data_format,
            input_data_format=input_data_format,
        )
        return padded_image

    
    def pad(
        self,
        images: List[np.ndarray],
        constant_values: Union[float, Iterable[float]] = 0,
        return_pixel_mask: bool = True,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: Optional[ChannelDimension] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> BatchFeature:
        
        pad_size = get_max_height_width(images, input_data_format=input_data_format)

        padded_images = [
            self._pad_image(
                image,
                pad_size,
                constant_values=constant_values,
                data_format=data_format,
                input_data_format=input_data_format,
            )
            for image in images
        ]
        data = {"pixel_values": padded_images}

        if return_pixel_mask:
            masks = [
                make_pixel_mask(image=image, output_size=pad_size, input_data_format=input_data_format)
                for image in images
            ]
            data["pixel_mask"] = masks

        return BatchFeature(data=data, tensor_type=return_tensors)

    def encode_inputs(
        self,
        pixel_values_list: List[ImageInput],
        segmentation_maps: ImageInput = None,
        instance_id_to_semantic_id: Optional[Union[List[Dict[int, int]], Dict[int, int]]] = None,
        ignore_index: Optional[int] = None,
        reduce_labels: bool = False,
        return_tensors: Optional[Union[str, TensorType]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ):
        
        ignore_index = self.ignore_index if ignore_index is None else ignore_index
        reduce_labels = self.reduce_labels if reduce_labels is None else reduce_labels

        pixel_values_list = [to_numpy_array(pixel_values) for pixel_values in pixel_values_list]

        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(pixel_values_list[0])

        encoded_inputs = self.pad(
            pixel_values_list, return_tensors=return_tensors, input_data_format=input_data_format
        )

        if segmentation_maps is not None:
            mask_labels = []
            class_labels = []
            pad_size = get_max_height_width(pixel_values_list)
            
            for idx, segmentation_map in enumerate(segmentation_maps):
                segmentation_map = to_numpy_array(segmentation_map)
                if isinstance(instance_id_to_semantic_id, list):
                    instance_id = instance_id_to_semantic_id[idx]
                else:
                    instance_id = instance_id_to_semantic_id
                
                masks, classes = self.convert_segmentation_map_to_binary_masks(
                    segmentation_map, instance_id, ignore_index=ignore_index, reduce_labels=reduce_labels
                )
                
                
                masks = [mask[None, ...] for mask in masks]
                masks = [
                    self._pad_image(image=mask, output_size=pad_size, constant_values=ignore_index) for mask in masks
                ]
                masks = np.concatenate(masks, axis=0)
                mask_labels.append(torch.from_numpy(masks))
                class_labels.append(torch.from_numpy(classes))

            
            encoded_inputs["mask_labels"] = mask_labels
            encoded_inputs["class_labels"] = class_labels

        return encoded_inputs

    def post_process_semantic_segmentation(
        self, outputs, target_sizes: Optional[List[Tuple[int, int]]] = None
    ) -> "torch.Tensor":
        
        class_queries_logits = outputs.class_queries_logits  
        masks_queries_logits = outputs.masks_queries_logits  

        
        masks_queries_logits = torch.nn.functional.interpolate(
            masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False
        )

        
        masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
        masks_probs = masks_queries_logits.sigmoid()  

        
        segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
        batch_size = class_queries_logits.shape[0]

        
        if target_sizes is not None:
            if batch_size != len(target_sizes):
                raise ValueError(
                    "Make sure that you pass in as many target sizes as the batch dimension of the logits"
                )

            semantic_segmentation = []
            for idx in range(batch_size):
                resized_logits = torch.nn.functional.interpolate(
                    segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False
                )
                semantic_map = resized_logits[0].argmax(dim=0)
                semantic_segmentation.append(semantic_map)
        else:
            semantic_segmentation = segmentation.argmax(dim=1)
            semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])]

        return semantic_segmentation

    def post_process_instance_segmentation(
        self,
        outputs,
        threshold: float = 0.5,
        mask_threshold: float = 0.5,
        overlap_mask_area_threshold: float = 0.8,
        target_sizes: Optional[List[Tuple[int, int]]] = None,
        return_coco_annotation: Optional[bool] = False,
        return_binary_maps: Optional[bool] = False,
    ) -> List[Dict]:
        
        if return_coco_annotation and return_binary_maps:
            raise ValueError("return_coco_annotation and return_binary_maps can not be both set to True.")

        
        class_queries_logits = outputs.class_queries_logits
        
        masks_queries_logits = outputs.masks_queries_logits

        
        masks_queries_logits = torch.nn.functional.interpolate(
            masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False
        )

        device = masks_queries_logits.device
        num_classes = class_queries_logits.shape[-1] - 1
        num_queries = class_queries_logits.shape[-2]

        
        results: List[Dict[str, TensorType]] = []

        for i in range(class_queries_logits.shape[0]):
            mask_pred = masks_queries_logits[i]
            mask_cls = class_queries_logits[i]

            scores = torch.nn.functional.softmax(mask_cls, dim=-1)[:, :-1]
            labels = torch.arange(num_classes, device=device).unsqueeze(0).repeat(num_queries, 1).flatten(0, 1)

            scores_per_image, topk_indices = scores.flatten(0, 1).topk(num_queries, sorted=False)
            labels_per_image = labels[topk_indices]

            topk_indices = torch.div(topk_indices, num_classes, rounding_mode="floor")
            mask_pred = mask_pred[topk_indices]
            pred_masks = (mask_pred > 0).float()

            
            mask_scores_per_image = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / (
                pred_masks.flatten(1).sum(1) + 1e-6
            )
            pred_scores = scores_per_image * mask_scores_per_image
            pred_classes = labels_per_image

            segmentation = torch.zeros((384, 384)) - 1
            if target_sizes is not None:
                segmentation = torch.zeros(target_sizes[i]) - 1
                pred_masks = torch.nn.functional.interpolate(
                    pred_masks.unsqueeze(0), size=target_sizes[i], mode="nearest"
                )[0]

            instance_maps, segments = [], []
            current_segment_id = 0
            for j in range(num_queries):
                score = pred_scores[j].item()

                if not torch.all(pred_masks[j] == 0) and score >= threshold:
                    segmentation[pred_masks[j] == 1] = current_segment_id
                    segments.append(
                        {
                            "id": current_segment_id,
                            "label_id": pred_classes[j].item(),
                            "was_fused": False,
                            "score": round(score, 6),
                        }
                    )
                    current_segment_id += 1
                    instance_maps.append(pred_masks[j])

            
            if return_coco_annotation:
                segmentation = convert_segmentation_to_rle(segmentation)

            
            if return_binary_maps and len(instance_maps) != 0:
                segmentation = torch.stack(instance_maps, dim=0)

            results.append({"segmentation": segmentation, "segments_info": segments})
        return results

    def post_process_panoptic_segmentation(
        self,
        outputs,
        threshold: float = 0.5,
        mask_threshold: float = 0.5,
        overlap_mask_area_threshold: float = 0.8,
        label_ids_to_fuse: Optional[Set[int]] = None,
        target_sizes: Optional[List[Tuple[int, int]]] = None,
    ) -> List[Dict]:
        

        if label_ids_to_fuse is None:
            logger.warning("`label_ids_to_fuse` unset. No instance will be fused.")
            label_ids_to_fuse = set()

        class_queries_logits = outputs.class_queries_logits  
        masks_queries_logits = outputs.masks_queries_logits  

        
        masks_queries_logits = torch.nn.functional.interpolate(
            masks_queries_logits, size=(384, 384), mode="bilinear", align_corners=False
        )

        batch_size = class_queries_logits.shape[0]
        num_labels = class_queries_logits.shape[-1] - 1

        mask_probs = masks_queries_logits.sigmoid()  

        
        pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1)

        
        results: List[Dict[str, TensorType]] = []

        for i in range(batch_size):
            mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects(
                mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels
            )

            
            if mask_probs_item.shape[0] <= 0:
                height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:]
                segmentation = torch.zeros((height, width)) - 1
                results.append({"segmentation": segmentation, "segments_info": []})
                continue

            
            target_size = target_sizes[i] if target_sizes is not None else None
            segmentation, segments = compute_segments(
                mask_probs=mask_probs_item,
                pred_scores=pred_scores_item,
                pred_labels=pred_labels_item,
                mask_threshold=mask_threshold,
                overlap_mask_area_threshold=overlap_mask_area_threshold,
                label_ids_to_fuse=label_ids_to_fuse,
                target_size=target_size,
            )

            results.append({"segmentation": segmentation, "segments_info": segments})
        return results