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 flava import math import random from functools import lrucache from typing import any dict iterable list optional tuple union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import resize tochanneldimensionformat from imageutils import openaiclipmean openaiclipstd channeldimension imageinput pilimageresampling inferchanneldimensionformat isscaledimage makelistofimages tonumpyarray validimages from utils import tensortype isvisionavailable logging if isvisionavailable import pil logger logging getloggername these values are taken from clip flavaimagemean openaiclipmean flavaimagestd openaiclipstd flavacodebookmean 0 0 0 0 0 0 flavacodebookstd 1 0 1 0 1 0 logitlaplaceeps float 0 1 inspired from https github commicrosoftunilmblobmasterbeitmaskinggenerator py class flavamaskinggenerator def init self inputsize unionint tupleint int 14 totalmaskpatches int 75 maskgroupmaxpatches optionalint none maskgroupminpatches int 16 maskgroupminaspectratio optionalfloat 0 3 maskgroupmaxaspectratio float none if not isinstanceinputsize tuple inputsize inputsize 2 self height self width inputsize self numpatches self height self width self totalmaskpatches totalmaskpatches self maskgroupminpatches maskgroupminpatches self maskgroupmaxpatches totalmaskpatches if maskgroupmaxpatches is none else maskgroupmaxpatches maskgroupmaxaspectratio maskgroupmaxaspectratio or 1 maskgroupminaspectratio self logaspectratio math logmaskgroupminaspectratio math logmaskgroupmaxaspectratio def reprself reprstr maskinggeneratord d d d max d 3f 3f self height self width self maskgroupminpatches self maskgroupmaxpatches self totalmaskpatches self logaspectratio0 self logaspectratio1 return reprstr def getshapeself return self height self width def maskself mask maxmaskpatches delta 0 for attempt in range10 targetarea random uniformself maskgroupminpatches maxmaskpatches aspectratio math exprandom uniformself logaspectratio height introundmath sqrttargetarea aspectratio width introundmath sqrttargetarea aspectratio if width self width and height self height top random randint0 self height height left random randint0 self width width nummasked masktop top height left left width sum overlap if 0 height width nummasked maxmaskpatches for i in rangetop top height for j in rangeleft left width if maski j 0 maski j 1 delta 1 if delta 0 break return delta def callself mask np zerosshapeself getshape dtypeint maskcount 0 while maskcount self totalmaskpatches maxmaskpatches self totalmaskpatches maskcount maxmaskpatches minmaxmaskpatches self maskgroupmaxpatches delta self maskmask maxmaskpatches if delta 0 break else maskcount delta return mask class flavaimageprocessorbaseimageprocessor r constructs a flava image processor args doresize bool optional defaults to true whether to resize the image s height width dimensions to the specified size can be overridden by the doresize parameter in preprocess size dictstr int optional defaults to height 224 width 224 size of the image after resizing can be overridden by the size parameter in preprocess resample pilimageresampling optional defaults to pilimageresampling bicubic resampling filter to use if resizing the image can be overridden by the resample parameter in preprocess docentercrop bool optional defaults to true whether to center crop the images can be overridden by the docentercrop parameter in preprocess cropsize dictstr int optional defaults to height 224 width 224 size of image after the center crop cropsizeheight cropsizewidth can be overridden by the cropsize parameter in preprocess dorescale bool optional defaults to true whether to rescale the image by the specified scale rescalefactor can be overridden by the dorescale parameter in preprocess rescalefactor int or float optional defaults to 1255 scale factor to use if rescaling the image can be overridden by the rescalefactor parameter in preprocess donormalize bool optional defaults to true whether to normalize the image can be overridden by the donormalize parameter in preprocess imagemean float or listfloat optional defaults to imagenetstandardmean mean to use if normalizing the image this is a float or list of floats the length of the number of channels in the image can be overridden by the imagemean parameter in the preprocess method imagestd float or listfloat optional defaults to imagenetstandardstd standard deviation to use if normalizing the image this is a float or list of floats the length of the number of channels in the image can be overridden by the imagestd parameter in the preprocess method returnimagemask bool optional defaults to false whether to return the image mask can be overridden by the returnimagemask parameter in preprocess inputsizepatches int optional defaults to 14 number of patches in the image in height and width direction 14x14 196 total patches can be overridden by the inputsizepatches parameter in preprocess totalmaskpatches int optional defaults to 75 total number of patches that should be masked can be overridden by the totalmaskpatches parameter in preprocess maskgroupminpatches int optional defaults to 16 minimum number of patches that should be masked can be overridden by the maskgroupminpatches parameter in preprocess maskgroupmaxpatches int optional maximum number of patches that should be masked can be overridden by the maskgroupmaxpatches parameter in preprocess maskgroupminaspectratio float optional defaults to 0 3 minimum aspect ratio of the mask window can be overridden by the maskgroupminaspectratio parameter in preprocess maskgroupmaxaspectratio float optional maximum aspect ratio of the mask window can be overridden by the maskgroupmaxaspectratio parameter in preprocess codebookdoresize bool optional defaults to true whether to resize the input for codebook to a certain can be overridden by the codebookdoresize parameter in preprocess codebooksize codebooksize dictstr int optional defaults to height 224 width 224 resize the input for codebook to the given size can be overridden by the codebooksize parameter in preprocess codebookresample pilimageresampling optional defaults to pilimageresampling lanczos resampling filter to use if resizing the codebook image can be overridden by the codebookresample parameter in preprocess codebookdocentercrop bool optional defaults to true whether to crop the input for codebook at the center if the input size is smaller than codebookcropsize along any edge the image is padded with 0 s and then center cropped can be overridden by the codebookdocentercrop parameter in preprocess codebookcropsize dictstr int optional defaults to height 224 width 224 desired output size for codebook input when applying centercropping can be overridden by the codebookcropsize parameter in preprocess codebookdorescale bool optional defaults to true whether to rescale the input for codebook by the specified scale codebookrescalefactor can be overridden by the codebookdorescale parameter in preprocess codebookrescalefactor int or float optional defaults to 1255 defines the scale factor to use if rescaling the codebook image can be overridden by the codebookrescalefactor parameter in preprocess codebookdomappixels bool optional defaults to true whether to map the pixel values of the codebook input to 1 2ex e can be overridden by the codebookdomappixels parameter in preprocess codebookdonormalize bool optional defaults to true whether or not to normalize the input for codebook with codebookimagemean and codebookimagestd can be overridden by the codebookdonormalize parameter in preprocess codebookimagemean optionalunionfloat iterablefloat optional defaults to 0 0 0 the sequence of means for each channel to be used when normalizing images for codebook can be overridden by the codebookimagemean parameter in preprocess codebookimagestd optionalunionfloat iterablefloat optional defaults to 0 5 0 5 0 5 the sequence of standard deviations for each channel to be used when normalizing images for codebook can be overridden by the codebookimagestd parameter in preprocess mask related params codebook related params 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 flavaimageprocessor frompretrainedcheckpoint codebooksize600 copied from transformers models vit imageprocessingvit vitimageprocessor resize with pilimageresampling bilinearpilimageresampling bicubic resize an image to sizeheight sizewidth args image np ndarray image to resize size dictstr int dictionary in the format height int width int specifying the size of the output image resample pilimageresampling optional defaults to pilimageresampling bicubic pilimageresampling filter to use when resizing the image e g pilimageresampling bicubic 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 can be one of channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format none or channeldimension none image in height width format inputdataformat channeldimension or str optional the channel dimension format for the input image if unset the channel dimension format 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 none or channeldimension none image in height width format returns np ndarray the resized image preprocesses a single image if doresize and size is none or resample is none raise valueerrorsize and resample must be specified if doresize is true if dorescale and rescalefactor is none raise valueerrorrescale factor must be specified if dorescale is true if donormalize and imagemean is none or imagestd is none raise valueerrorimage mean and std must be specified if donormalize is true 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 we assume that all images have the same channel dimension format inputdataformat inferchanneldimensionformatimage if doresize image self resizeimageimage sizesize resampleresample inputdataformatinputdataformat if docentercrop image self centercropimageimage sizecropsize inputdataformatinputdataformat if dorescale image self rescaleimageimage scalerescalefactor inputdataformatinputdataformat if donormalize image self normalizeimageimage meanimagemean stdimagestd inputdataformatinputdataformat if domappixels image self mappixelsimage if dataformat is not none image tochanneldimensionformatimage dataformat inputchanneldiminputdataformat return image def preprocess self images imageinput doresize optionalbool none size dictstr int none resample pilimageresampling none docentercrop optionalbool none cropsize optionaldictstr int none dorescale optionalbool none rescalefactor optionalfloat none donormalize optionalbool none imagemean optionalunionfloat listfloat none imagestd optionalunionfloat listfloat none mask related params returnimagemask optionalbool none inputsizepatches optionalint none totalmaskpatches optionalint none maskgroupminpatches optionalint none maskgroupmaxpatches optionalint none maskgroupminaspectratio optionalfloat none maskgroupmaxaspectratio optionalfloat none codebook related params returncodebookpixels optionalbool none codebookdoresize optionalbool none codebooksize optionaldictstr int none codebookresample optionalint none codebookdocentercrop optionalbool none codebookcropsize optionaldictstr int none codebookdorescale optionalbool none codebookrescalefactor optionalfloat none codebookdomappixels optionalbool none codebookdonormalize optionalbool none codebookimagemean optionaliterablefloat none codebookimagestd optionaliterablefloat none returntensors optionalunionstr tensortype none dataformat channeldimension channeldimension first inputdataformat optionalunionstr channeldimension none kwargs pil image image doresize doresize if doresize is not none else self doresize size size if size is not none else self size size getsizedictsize resample resample if resample is not none else self resample docentercrop docentercrop if docentercrop is not none else self docentercrop cropsize cropsize if cropsize is not none else self cropsize cropsize getsizedictcropsize paramnamecropsize dorescale dorescale if dorescale is not none else self dorescale rescalefactor rescalefactor if rescalefactor is not none else self rescalefactor donormalize donormalize if donormalize is not none else self donormalize imagemean imagemean if imagemean is not none else self imagemean imagestd imagestd if imagestd is not none else self imagestd returnimagemask returnimagemask if returnimagemask is not none else self returnimagemask inputsizepatches inputsizepatches if inputsizepatches is not none else self inputsizepatches totalmaskpatches totalmaskpatches if totalmaskpatches is not none else self totalmaskpatches maskgroupminpatches maskgroupminpatches if maskgroupminpatches is not none else self maskgroupminpatches maskgroupmaxpatches maskgroupmaxpatches if maskgroupmaxpatches is not none else self maskgroupmaxpatches maskgroupminaspectratio maskgroupminaspectratio if maskgroupminaspectratio is not none else self maskgroupminaspectratio maskgroupmaxaspectratio maskgroupmaxaspectratio if maskgroupmaxaspectratio is not none else self maskgroupmaxaspectratio returncodebookpixels returncodebookpixels if returncodebookpixels is not none else self returncodebookpixels codebookdoresize codebookdoresize if codebookdoresize is not none else self codebookdoresize codebooksize codebooksize if codebooksize is not none else self codebooksize codebooksize getsizedictcodebooksize paramnamecodebooksize codebookresample codebookresample if codebookresample is not none else self codebookresample codebookdorescale codebookdorescale if codebookdorescale is not none else self codebookdorescale codebookrescalefactor codebookrescalefactor if codebookrescalefactor is not none else self codebookrescalefactor codebookdocentercrop codebookdocentercrop if codebookdocentercrop is not none else self codebookdocentercrop codebookcropsize codebookcropsize if codebookcropsize is not none else self codebookcropsize codebookcropsize getsizedictcodebookcropsize paramnamecodebookcropsize codebookdomappixels codebookdomappixels if codebookdomappixels is not none else self codebookdomappixels codebookdonormalize codebookdonormalize if codebookdonormalize is not none else self codebookdonormalize codebookimagemean codebookimagemean if codebookimagemean is not none else self codebookimagemean codebookimagestd codebookimagestd if codebookimagestd is not none else self codebookimagestd images makelistofimagesimages if not validimagesimages raise valueerror invalid image type must be of type pil image image numpy ndarray torch tensor tf tensor or jax ndarray processedimages self preprocessimage imageimg doresizedoresize sizesize resampleresample docentercropdocentercrop cropsizecropsize dorescaledorescale rescalefactorrescalefactor donormalizedonormalize imagemeanimagemean imagestdimagestd domappixelsfalse dataformatdataformat inputdataformatinputdataformat for img in images data pixelvalues processedimages if returncodebookpixels codebookimages self preprocessimage imageimg doresizecodebookdoresize sizecodebooksize resamplecodebookresample docentercropcodebookdocentercrop cropsizecodebookcropsize dorescalecodebookdorescale rescalefactorcodebookrescalefactor donormalizecodebookdonormalize imagemeancodebookimagemean imagestdcodebookimagestd domappixelscodebookdomappixels dataformatdataformat inputdataformatinputdataformat for img in images datacodebookpixelvalues codebookimages if returnimagemask maskgenerator self maskinggenerator inputsizepatchesinputsizepatches totalmaskpatchestotalmaskpatches maskgroupminpatchesmaskgroupminpatches maskgroupmaxpatchesmaskgroupmaxpatches maskgroupminaspectratiomaskgroupminaspectratio maskgroupmaxaspectratiomaskgroupmaxaspectratio masks maskgenerator for in images databoolmaskedpos masks return batchfeaturedatadata tensortypereturntensors 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 flava these values are taken from clip inspired from https github com microsoft unilm blob master beit masking_generator py overlap constructs a flava image processor args do_resize bool optional defaults to true whether to resize the image s height width dimensions to the specified size can be overridden by the do_resize parameter in preprocess size dict str int optional defaults to height 224 width 224 size of the image after resizing can be overridden by the size parameter in preprocess resample pilimageresampling optional defaults to pilimageresampling bicubic resampling filter to use if resizing the image can be overridden by the resample parameter in preprocess do_center_crop bool optional defaults to true whether to center crop the images can be overridden by the do_center_crop parameter in preprocess crop_size dict str int optional defaults to height 224 width 224 size of image after the center crop crop_size height crop_size width can be overridden by the crop_size parameter in preprocess do_rescale bool optional defaults to true whether to rescale the image by the specified scale rescale_factor can be overridden by the do_rescale parameter in preprocess rescale_factor int or float optional defaults to 1 255 scale factor to use if rescaling the image can be overridden by the rescale_factor parameter in preprocess do_normalize bool optional defaults to true whether to normalize the image can be overridden by the do_normalize parameter in preprocess image_mean float or list float optional defaults to imagenet_standard_mean mean to use if normalizing the image this is a float or list of floats the length of the number of channels in the image can be overridden by the image_mean parameter in the preprocess method image_std float or list float optional defaults to imagenet_standard_std standard deviation to use if normalizing the image this is a float or list of floats the length of the number of channels in the image can be overridden by the image_std parameter in the preprocess method return_image_mask bool optional defaults to false whether to return the image mask can be overridden by the return_image_mask parameter in preprocess input_size_patches int optional defaults to 14 number of patches in the image in height and width direction 14x14 196 total patches can be overridden by the input_size_patches parameter in preprocess total_mask_patches int optional defaults to 75 total number of patches that should be masked can be overridden by the total_mask_patches parameter in preprocess mask_group_min_patches int optional defaults to 16 minimum number of patches that should be masked can be overridden by the mask_group_min_patches parameter in preprocess mask_group_max_patches int optional maximum number of patches that should be masked can be overridden by the mask_group_max_patches parameter in preprocess mask_group_min_aspect_ratio float optional defaults to 0 3 minimum aspect ratio of the mask window can be overridden by the mask_group_min_aspect_ratio parameter in preprocess mask_group_max_aspect_ratio float optional maximum aspect ratio of the mask window can be overridden by the mask_group_max_aspect_ratio parameter in preprocess codebook_do_resize bool optional defaults to true whether to resize the input for codebook to a certain can be overridden by the codebook_do_resize parameter in preprocess codebook_size codebook_size dict str int optional defaults to height 224 width 224 resize the input for codebook to the given size can be overridden by the codebook_size parameter in preprocess codebook_resample pilimageresampling optional defaults to pilimageresampling lanczos resampling filter to use if resizing the codebook image can be overridden by the codebook_resample parameter in preprocess codebook_do_center_crop bool optional defaults to true whether to crop the input for codebook at the center if the input size is smaller than codebook_crop_size along any edge the image is padded with 0 s and then center cropped can be overridden by the codebook_do_center_crop parameter in preprocess codebook_crop_size dict str int optional defaults to height 224 width 224 desired output size for codebook input when applying center cropping can be overridden by the codebook_crop_size parameter in preprocess codebook_do_rescale bool optional defaults to true whether to rescale the input for codebook by the specified scale codebook_rescale_factor can be overridden by the codebook_do_rescale parameter in preprocess codebook_rescale_factor int or float optional defaults to 1 255 defines the scale factor to use if rescaling the codebook image can be overridden by the codebook_rescale_factor parameter in preprocess codebook_do_map_pixels bool optional defaults to true whether to map the pixel values of the codebook input to 1 2e x e can be overridden by the codebook_do_map_pixels parameter in preprocess codebook_do_normalize bool optional defaults to true whether or not to normalize the input for codebook with codebook_image_mean and codebook_image_std can be overridden by the codebook_do_normalize parameter in preprocess codebook_image_mean optional union float iterable float optional defaults to 0 0 0 the sequence of means for each channel to be used when normalizing images for codebook can be overridden by the codebook_image_mean parameter in preprocess codebook_image_std optional union float iterable float optional defaults to 0 5 0 5 0 5 the sequence of standard deviations for each channel to be used when normalizing images for codebook can be overridden by the codebook_image_std parameter in preprocess mask related params codebook related params 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 flavaimageprocessor from_pretrained checkpoint codebook_size 600 copied from transformers models vit image_processing_vit vitimageprocessor resize with pilimageresampling bilinear pilimageresampling bicubic resize an image to size height size width args image np ndarray image to resize size dict str int dictionary in the format height int width int specifying the size of the output image resample pilimageresampling optional defaults to pilimageresampling bicubic pilimageresampling filter to use when resizing the image e g pilimageresampling bicubic 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 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 none or channeldimension none image in height width format input_data_format channeldimension or str optional the channel dimension format for the input image if unset the channel dimension format 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 none or channeldimension none image in height width format returns np ndarray the resized image preprocesses a single image all transformations expect numpy arrays we assume that all images have the same channel dimension format mask related params codebook related params preprocess an image or batch of images args images imageinput image to preprocess expects a single or batch of images with pixel values ranging from 0 to 255 if passing in images with pixel values between 0 and 1 set do_rescale false do_resize bool optional defaults to self do_resize whether to resize the image size dict str int optional defaults to self size size of the image resample int optional defaults to self resample resampling filter to use if resizing the image this can be one of the enum pilimageresampling only has an effect if do_resize is set to true do_center_crop bool optional defaults to self do_center_crop whether to center crop the image crop_size dict str int optional defaults to self crop_size size of the center crop only has an effect if do_center_crop is set to true do_rescale bool optional defaults to self do_rescale whether to rescale the image values between 0 1 rescale_factor float optional defaults to self rescale_factor rescale factor to rescale the image by if do_rescale is set to true do_normalize bool optional defaults to self do_normalize whether to normalize the image image_mean float or list float optional defaults to self image_mean image mean image_std float or list float optional defaults to self image_std image standard deviation return_image_mask bool optional defaults to self return_image_mask whether to return the image mask input_size_patches int optional defaults to self input_size_patches size of the patches to extract from the image total_mask_patches int optional defaults to self total_mask_patches total number of patches to extract from the image mask_group_min_patches int optional defaults to self mask_group_min_patches minimum number of patches to extract from the image mask_group_max_patches int optional defaults to self mask_group_max_patches maximum number of patches to extract from the image mask_group_min_aspect_ratio float optional defaults to self mask_group_min_aspect_ratio minimum aspect ratio of the patches to extract from the image mask_group_max_aspect_ratio float optional defaults to self mask_group_max_aspect_ratio maximum aspect ratio of the patches to extract from the image return_codebook_pixels bool optional defaults to self return_codebook_pixels whether to return the codebook pixels codebook_do_resize bool optional defaults to self codebook_do_resize whether to resize the codebook pixels codebook_size dict str int optional defaults to self codebook_size size of the codebook pixels codebook_resample int optional defaults to self codebook_resample resampling filter to use if resizing the codebook pixels this can be one of the enum pilimageresampling only has an effect if codebook_do_resize is set to true codebook_do_center_crop bool optional defaults to self codebook_do_center_crop whether to center crop the codebook pixels codebook_crop_size dict str int optional defaults to self codebook_crop_size size of the center crop of the codebook pixels only has an effect if codebook_do_center_crop is set to true codebook_do_rescale bool optional defaults to self codebook_do_rescale whether to rescale the codebook pixels values between 0 1 codebook_rescale_factor float optional defaults to self codebook_rescale_factor rescale factor to rescale the codebook pixels by if codebook_do_rescale is set to true codebook_do_map_pixels bool optional defaults to self codebook_do_map_pixels whether to map the codebook pixels values codebook_do_normalize bool optional defaults to self codebook_do_normalize whether to normalize the codebook pixels codebook_image_mean float or list float optional defaults to self codebook_image_mean codebook pixels mean to normalize the codebook pixels by if codebook_do_normalize is set to true codebook_image_std float or list float optional defaults to self codebook_image_std codebook pixels standard deviation to normalize the codebook pixels by if codebook_do_normalize is set to true 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 channeldimension or str optional defaults to channeldimension first the channel dimension format for the output image can be one of channeldimension first image in num_channels height width format channeldimension last image in height width num_channels format input_data_format channeldimension or str optional the channel dimension format for the input image if unset the channel dimension format 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 none or channeldimension none image in height width format
import math
import random
from functools import lru_cache
from typing import Any, Dict, Iterable, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import resize, to_channel_dimension_format
from ...image_utils import (
    OPENAI_CLIP_MEAN,
    OPENAI_CLIP_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    infer_channel_dimension_format,
    is_scaled_image,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
    import PIL


logger = logging.get_logger(__name__)



FLAVA_IMAGE_MEAN = OPENAI_CLIP_MEAN
FLAVA_IMAGE_STD = OPENAI_CLIP_STD
FLAVA_CODEBOOK_MEAN = [0.0, 0.0, 0.0]
FLAVA_CODEBOOK_STD = [1.0, 1.0, 1.0]
LOGIT_LAPLACE_EPS: float = 0.1



class FlavaMaskingGenerator:
    def __init__(
        self,
        input_size: Union[int, Tuple[int, int]] = 14,
        total_mask_patches: int = 75,
        mask_group_max_patches: Optional[int] = None,
        mask_group_min_patches: int = 16,
        mask_group_min_aspect_ratio: Optional[float] = 0.3,
        mask_group_max_aspect_ratio: float = None,
    ):
        if not isinstance(input_size, tuple):
            input_size = (input_size,) * 2
        self.height, self.width = input_size

        self.num_patches = self.height * self.width
        self.total_mask_patches = total_mask_patches

        self.mask_group_min_patches = mask_group_min_patches
        self.mask_group_max_patches = total_mask_patches if mask_group_max_patches is None else mask_group_max_patches

        mask_group_max_aspect_ratio = mask_group_max_aspect_ratio or 1 / mask_group_min_aspect_ratio
        self.log_aspect_ratio = (math.log(mask_group_min_aspect_ratio), math.log(mask_group_max_aspect_ratio))

    def __repr__(self):
        repr_str = "MaskingGenerator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
            self.height,
            self.width,
            self.mask_group_min_patches,
            self.mask_group_max_patches,
            self.total_mask_patches,
            self.log_aspect_ratio[0],
            self.log_aspect_ratio[1],
        )
        return repr_str

    def get_shape(self):
        return self.height, self.width

    def _mask(self, mask, max_mask_patches):
        delta = 0
        for _attempt in range(10):
            target_area = random.uniform(self.mask_group_min_patches, max_mask_patches)
            aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
            height = int(round(math.sqrt(target_area * aspect_ratio)))
            width = int(round(math.sqrt(target_area / aspect_ratio)))
            if width < self.width and height < self.height:
                top = random.randint(0, self.height - height)
                left = random.randint(0, self.width - width)

                num_masked = mask[top : top + height, left : left + width].sum()
                
                if 0 < height * width - num_masked <= max_mask_patches:
                    for i in range(top, top + height):
                        for j in range(left, left + width):
                            if mask[i, j] == 0:
                                mask[i, j] = 1
                                delta += 1

                if delta > 0:
                    break
        return delta

    def __call__(self):
        mask = np.zeros(shape=self.get_shape(), dtype=int)
        mask_count = 0
        while mask_count < self.total_mask_patches:
            max_mask_patches = self.total_mask_patches - mask_count
            max_mask_patches = min(max_mask_patches, self.mask_group_max_patches)

            delta = self._mask(mask, max_mask_patches)
            if delta == 0:
                break
            else:
                mask_count += delta

        return mask


class FlavaImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        do_center_crop: bool = True,
        crop_size: Dict[str, int] = None,
        do_rescale: bool = True,
        rescale_factor: Union[int, float] = 1 / 255,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, Iterable[float]]] = None,
        image_std: Optional[Union[float, Iterable[float]]] = None,
        
        return_image_mask: bool = False,
        input_size_patches: int = 14,
        total_mask_patches: int = 75,
        mask_group_min_patches: int = 16,
        mask_group_max_patches: Optional[int] = None,
        mask_group_min_aspect_ratio: float = 0.3,
        mask_group_max_aspect_ratio: Optional[float] = None,
        
        return_codebook_pixels: bool = False,
        codebook_do_resize: bool = True,
        codebook_size: bool = None,
        codebook_resample: int = PILImageResampling.LANCZOS,
        codebook_do_center_crop: bool = True,
        codebook_crop_size: int = None,
        codebook_do_rescale: bool = True,
        codebook_rescale_factor: Union[int, float] = 1 / 255,
        codebook_do_map_pixels: bool = True,
        codebook_do_normalize: bool = True,
        codebook_image_mean: Optional[Union[float, Iterable[float]]] = None,
        codebook_image_std: Optional[Union[float, Iterable[float]]] = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        size = size if size is not None else {"height": 224, "width": 224}
        size = get_size_dict(size)
        crop_size = crop_size if crop_size is not None else {"height": 224, "width": 224}
        crop_size = get_size_dict(crop_size, param_name="crop_size")

        codebook_size = codebook_size if codebook_size is not None else {"height": 112, "width": 112}
        codebook_size = get_size_dict(codebook_size, param_name="codebook_size")
        codebook_crop_size = codebook_crop_size if codebook_crop_size is not None else {"height": 112, "width": 112}
        codebook_crop_size = get_size_dict(codebook_crop_size, param_name="codebook_crop_size")

        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_center_crop = do_center_crop
        self.crop_size = crop_size
        self.do_normalize = do_normalize
        self.image_mean = image_mean if image_mean is not None else FLAVA_IMAGE_MEAN
        self.image_std = image_std if image_std is not None else FLAVA_IMAGE_STD

        self.return_image_mask = return_image_mask
        self.input_size_patches = input_size_patches
        self.total_mask_patches = total_mask_patches
        self.mask_group_min_patches = mask_group_min_patches
        self.mask_group_max_patches = mask_group_max_patches
        self.mask_group_min_aspect_ratio = mask_group_min_aspect_ratio
        self.mask_group_max_aspect_ratio = mask_group_max_aspect_ratio

        self.return_codebook_pixels = return_codebook_pixels
        self.codebook_do_resize = codebook_do_resize
        self.codebook_size = codebook_size
        self.codebook_resample = codebook_resample
        self.codebook_do_center_crop = codebook_do_center_crop
        self.codebook_crop_size = codebook_crop_size
        self.codebook_do_rescale = codebook_do_rescale
        self.codebook_rescale_factor = codebook_rescale_factor
        self.codebook_do_map_pixels = codebook_do_map_pixels
        self.codebook_do_normalize = codebook_do_normalize
        self.codebook_image_mean = codebook_image_mean
        self.codebook_image_mean = codebook_image_mean if codebook_image_mean is not None else FLAVA_CODEBOOK_MEAN
        self.codebook_image_std = codebook_image_std if codebook_image_std is not None else FLAVA_CODEBOOK_STD

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

    @lru_cache()
    def masking_generator(
        self,
        input_size_patches,
        total_mask_patches,
        mask_group_min_patches,
        mask_group_max_patches,
        mask_group_min_aspect_ratio,
        mask_group_max_aspect_ratio,
    ) -> FlavaMaskingGenerator:
        return FlavaMaskingGenerator(
            input_size=input_size_patches,
            total_mask_patches=total_mask_patches,
            mask_group_min_patches=mask_group_min_patches,
            mask_group_max_patches=mask_group_max_patches,
            mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
            mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
        )

    
    def resize(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> np.ndarray:
        
        size = get_size_dict(size)
        if "height" not in size or "width" not in size:
            raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}")
        output_size = (size["height"], size["width"])
        return resize(
            image,
            size=output_size,
            resample=resample,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )

    def map_pixels(self, image: np.ndarray) -> np.ndarray:
        return (1 - 2 * LOGIT_LAPLACE_EPS) * image + LOGIT_LAPLACE_EPS

    def _preprocess_image(
        self,
        image: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_center_crop: bool = None,
        crop_size: Dict[str, int] = 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,
        do_map_pixels: bool = None,
        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[ChannelDimension] = None,
    ) -> np.ndarray:
        
        if do_resize and size is None or resample is None:
            raise ValueError("Size and resample must be specified if do_resize is True.")

        if do_rescale and rescale_factor is None:
            raise ValueError("Rescale factor must be specified if do_rescale is True.")

        if do_normalize and (image_mean is None or image_std is None):
            raise ValueError("Image mean and std must be specified if do_normalize is True.")

        
        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)

        if do_resize:
            image = self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)

        if do_center_crop:
            image = self.center_crop(image=image, size=crop_size, input_data_format=input_data_format)

        if do_rescale:
            image = self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)

        if do_normalize:
            image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)

        if do_map_pixels:
            image = self.map_pixels(image)

        if data_format is not None:
            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
        return image

    def preprocess(
        self,
        images: ImageInput,
        do_resize: Optional[bool] = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_center_crop: Optional[bool] = None,
        crop_size: Optional[Dict[str, int]] = 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,
        
        return_image_mask: Optional[bool] = None,
        input_size_patches: Optional[int] = None,
        total_mask_patches: Optional[int] = None,
        mask_group_min_patches: Optional[int] = None,
        mask_group_max_patches: Optional[int] = None,
        mask_group_min_aspect_ratio: Optional[float] = None,
        mask_group_max_aspect_ratio: Optional[float] = None,
        
        return_codebook_pixels: Optional[bool] = None,
        codebook_do_resize: Optional[bool] = None,
        codebook_size: Optional[Dict[str, int]] = None,
        codebook_resample: Optional[int] = None,
        codebook_do_center_crop: Optional[bool] = None,
        codebook_crop_size: Optional[Dict[str, int]] = None,
        codebook_do_rescale: Optional[bool] = None,
        codebook_rescale_factor: Optional[float] = None,
        codebook_do_map_pixels: Optional[bool] = None,
        codebook_do_normalize: Optional[bool] = None,
        codebook_image_mean: Optional[Iterable[float]] = None,
        codebook_image_std: Optional[Iterable[float]] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: ChannelDimension = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> PIL.Image.Image:
        
        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)
        resample = resample if resample is not None else self.resample
        do_center_crop = do_center_crop if do_center_crop is not None else self.do_center_crop
        crop_size = crop_size if crop_size is not None else self.crop_size
        crop_size = get_size_dict(crop_size, param_name="crop_size")
        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

        return_image_mask = return_image_mask if return_image_mask is not None else self.return_image_mask
        input_size_patches = input_size_patches if input_size_patches is not None else self.input_size_patches
        total_mask_patches = total_mask_patches if total_mask_patches is not None else self.total_mask_patches
        mask_group_min_patches = (
            mask_group_min_patches if mask_group_min_patches is not None else self.mask_group_min_patches
        )
        mask_group_max_patches = (
            mask_group_max_patches if mask_group_max_patches is not None else self.mask_group_max_patches
        )
        mask_group_min_aspect_ratio = (
            mask_group_min_aspect_ratio
            if mask_group_min_aspect_ratio is not None
            else self.mask_group_min_aspect_ratio
        )
        mask_group_max_aspect_ratio = (
            mask_group_max_aspect_ratio
            if mask_group_max_aspect_ratio is not None
            else self.mask_group_max_aspect_ratio
        )

        return_codebook_pixels = (
            return_codebook_pixels if return_codebook_pixels is not None else self.return_codebook_pixels
        )
        codebook_do_resize = codebook_do_resize if codebook_do_resize is not None else self.codebook_do_resize
        codebook_size = codebook_size if codebook_size is not None else self.codebook_size
        codebook_size = get_size_dict(codebook_size, param_name="codebook_size")
        codebook_resample = codebook_resample if codebook_resample is not None else self.codebook_resample
        codebook_do_rescale = codebook_do_rescale if codebook_do_rescale is not None else self.codebook_do_rescale
        codebook_rescale_factor = (
            codebook_rescale_factor if codebook_rescale_factor is not None else self.codebook_rescale_factor
        )
        codebook_do_center_crop = (
            codebook_do_center_crop if codebook_do_center_crop is not None else self.codebook_do_center_crop
        )
        codebook_crop_size = codebook_crop_size if codebook_crop_size is not None else self.codebook_crop_size
        codebook_crop_size = get_size_dict(codebook_crop_size, param_name="codebook_crop_size")
        codebook_do_map_pixels = (
            codebook_do_map_pixels if codebook_do_map_pixels is not None else self.codebook_do_map_pixels
        )
        codebook_do_normalize = (
            codebook_do_normalize if codebook_do_normalize is not None else self.codebook_do_normalize
        )
        codebook_image_mean = codebook_image_mean if codebook_image_mean is not None else self.codebook_image_mean
        codebook_image_std = codebook_image_std if codebook_image_std is not None else self.codebook_image_std

        images = make_list_of_images(images)

        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."
            )

        processed_images = [
            self._preprocess_image(
                image=img,
                do_resize=do_resize,
                size=size,
                resample=resample,
                do_center_crop=do_center_crop,
                crop_size=crop_size,
                do_rescale=do_rescale,
                rescale_factor=rescale_factor,
                do_normalize=do_normalize,
                image_mean=image_mean,
                image_std=image_std,
                do_map_pixels=False,
                data_format=data_format,
                input_data_format=input_data_format,
            )
            for img in images
        ]
        data = {"pixel_values": processed_images}

        if return_codebook_pixels:
            codebook_images = [
                self._preprocess_image(
                    image=img,
                    do_resize=codebook_do_resize,
                    size=codebook_size,
                    resample=codebook_resample,
                    do_center_crop=codebook_do_center_crop,
                    crop_size=codebook_crop_size,
                    do_rescale=codebook_do_rescale,
                    rescale_factor=codebook_rescale_factor,
                    do_normalize=codebook_do_normalize,
                    image_mean=codebook_image_mean,
                    image_std=codebook_image_std,
                    do_map_pixels=codebook_do_map_pixels,
                    data_format=data_format,
                    input_data_format=input_data_format,
                )
                for img in images
            ]
            data["codebook_pixel_values"] = codebook_images

        if return_image_mask:
            mask_generator = self.masking_generator(
                input_size_patches=input_size_patches,
                total_mask_patches=total_mask_patches,
                mask_group_min_patches=mask_group_min_patches,
                mask_group_max_patches=mask_group_max_patches,
                mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
                mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
            )
            masks = [mask_generator() for _ in images]
            data["bool_masked_pos"] = masks

        return BatchFeature(data=data, tensor_type=return_tensors)