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 vilt from typing import any dict iterable list optional tuple union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import paddingmode pad resize tochanneldimensionformat from imageutils import imagenetstandardmean imagenetstandardstd channeldimension imageinput pilimageresampling getimagesize inferchanneldimensionformat isscaledimage makelistofimages tonumpyarray validimages from utils import tensortype isvisionavailable logging if isvisionavailable import pil logger logging getloggername def maxacrossindicesvalues iterableany listany return maxvaluesi for valuesi in zipvalues 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 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 def getresizeoutputimagesize inputimage np ndarray shorter int 800 longer int 1333 sizedivisor int 32 inputdataformat optionalunionstr channeldimension none tupleint int inputheight inputwidth getimagesizeinputimage inputdataformat minsize maxsize shorter longer scale minsize mininputheight inputwidth if inputheight inputwidth newheight minsize newwidth scale inputwidth else newheight scale inputheight newwidth minsize if maxnewheight newwidth maxsize scale maxsize maxnewheight newwidth newheight scale newheight newwidth scale newwidth newheight newwidth intnewheight 0 5 intnewwidth 0 5 newheight newheight sizedivisor sizedivisor newwidth newwidth sizedivisor sizedivisor return newheight newwidth class viltimageprocessorbaseimageprocessor r constructs a vilt 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 the preprocess method size dictstr int optional defaults to shortestedge 384 resize the shorter side of the input to sizeshortestedge the longer side will be limited to under int1333 800 sizeshortestedge while preserving the aspect ratio only has an effect if doresize is set to true can be overridden by the size parameter in the preprocess method sizedivisor int optional defaults to 32 the size by which to make sure both the height and width can be divided only has an effect if doresize is set to true can be overridden by the sizedivisor parameter in the preprocess method resample pilimageresampling optional defaults to resampling bicubic resampling filter to use if resizing the image only has an effect if doresize is set to true can be overridden by the resample parameter in the preprocess method dorescale bool optional defaults to true wwhether to rescale the image by the specified scale rescalefactor can be overridden by the dorescale parameter in the preprocess method rescalefactor int or float optional defaults to 1255 scale factor to use if rescaling the image only has an effect if dorescale is set to true can be overridden by the rescalefactor parameter in the preprocess method donormalize bool optional defaults to true whether to normalize the image can be overridden by the donormalize parameter in the preprocess method can be overridden by the donormalize parameter in the preprocess method 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 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 can be overridden by the imagestd parameter in the preprocess method dopad bool optional defaults to true whether to pad the image to the maxheight maxwidth of the images in the batch can be overridden by the dopad parameter in the preprocess method overrides the fromdict method from the base class to make sure reducelabels is updated if image processor is created using fromdict and kwargs e g viltimageprocessor frompretrainedcheckpoint padandreturnpixelmaskfalse resize an image resizes the shorter side of the image to sizeshortestedge while preserving the aspect ratio if the longer side is larger than the max size intsizeshortestedge 1333 800 the longer side is then resized to the max size while preserving the aspect ratio args image np ndarray image to resize size dictstr int controls the size of the output image should be of the form shortestedge int sizedivisor int defaults to 32 the image is resized to a size that is a multiple of this value resample pilimageresampling filter optional defaults to pilimageresampling bicubic resampling filter to use when resiizing the image 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 str or channeldimension optional the channel dimension format of the input image if not provided it will be inferred 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 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 dorescalefalse doresize bool optional defaults to self doresize whether to resize the image size dictstr int optional defaults to self size controls the size of the image after resize the shortest edge of the image is resized to sizeshortestedge whilst preserving the aspect ratio if the longest edge of this resized image is intsizeshortestedge 1333 800 then the image is resized again to make the longest edge equal to intsizeshortestedge 1333 800 sizedivisor int optional defaults to self sizedivisor the image is resized to a size that is a multiple of this value resample pilimageresampling optional defaults to self resample resampling filter to use if resizing the image only has an effect if doresize is set to true dorescale bool optional defaults to self dorescale whether to rescale the image values between 0 1 rescalefactor float optional defaults to self rescalefactor rescale factor to rescale the image by if dorescale is set to true donormalize bool optional defaults to self donormalize whether to normalize the image imagemean float or listfloat optional defaults to self imagemean image mean to normalize the image by if donormalize is set to true imagestd float or listfloat optional defaults to self imagestd image standard deviation to normalize the image by if donormalize is set to true dopad bool optional defaults to self dopad whether to pad the image to the maxheight maxwidth in the batch if true a pixel mask is also created and returned 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 channeldimension or str optional defaults to channeldimension first the channel dimension format for the output image can be one of channeldimension first image in numchannels height width format channeldimension last image in height width numchannels 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 all transformations expect numpy arrays we assume that all images have the same channel dimension format 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 vilt return the maximum value across all indices of an iterable of values 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 get the maximum height and width across all images in a batch constructs a vilt 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 the preprocess method size dict str int optional defaults to shortest_edge 384 resize the shorter side of the input to size shortest_edge the longer side will be limited to under int 1333 800 size shortest_edge while preserving the aspect ratio only has an effect if do_resize is set to true can be overridden by the size parameter in the preprocess method size_divisor int optional defaults to 32 the size by which to make sure both the height and width can be divided only has an effect if do_resize is set to true can be overridden by the size_divisor parameter in the preprocess method resample pilimageresampling optional defaults to resampling bicubic resampling filter to use if resizing the image only has an effect if do_resize is set to true can be overridden by the resample parameter in the preprocess method do_rescale bool optional defaults to true wwhether to rescale the image by the specified scale rescale_factor can be overridden by the do_rescale parameter in the preprocess method rescale_factor int or float optional defaults to 1 255 scale factor to use if rescaling the image only has an effect if do_rescale is set to true can be overridden by the rescale_factor parameter in the preprocess method do_normalize bool optional defaults to true whether to normalize the image can be overridden by the do_normalize parameter in the preprocess method can be overridden by the do_normalize parameter in the preprocess method 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 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 can be overridden by the image_std parameter in the preprocess method do_pad bool optional defaults to true whether to pad the image to the max_height max_width of the images in the batch can be overridden by the do_pad parameter in the preprocess method overrides the from_dict method from the base class to make sure reduce_labels is updated if image processor is created using from_dict and kwargs e g viltimageprocessor from_pretrained checkpoint pad_and_return_pixel_mask false resize an image resizes the shorter side of the image to size shortest_edge while preserving the aspect ratio if the longer side is larger than the max size int size shortest_edge 1333 800 the longer side is then resized to the max size while preserving the aspect ratio args image np ndarray image to resize size dict str int controls the size of the output image should be of the form shortest_edge int size_divisor int defaults to 32 the image is resized to a size that is a multiple of this value resample pilimageresampling filter optional defaults to pilimageresampling bicubic resampling filter to use when resiizing the image 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 str or channeldimension 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 _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 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 controls the size of the image after resize the shortest edge of the image is resized to size shortest_edge whilst preserving the aspect ratio if the longest edge of this resized image is int size shortest_edge 1333 800 then the image is resized again to make the longest edge equal to int size shortest_edge 1333 800 size_divisor int optional defaults to self size_divisor the image is resized to a size that is a multiple of this value resample pilimageresampling optional defaults to self resample resampling filter to use if resizing the image only has an effect if do_resize 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 to normalize the image by if do_normalize is set to true image_std float or list float optional defaults to self image_std image standard deviation to normalize the image by if do_normalize is set to true do_pad bool optional defaults to self do_pad whether to pad the image to the max_height max_width in the batch if true a pixel mask is also created and returned 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 all transformations expect numpy arrays we assume that all images have the same channel dimension format
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 PaddingMode, pad, resize, to_channel_dimension_format
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    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__)


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


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 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 get_resize_output_image_size(
    input_image: np.ndarray,
    shorter: int = 800,
    longer: int = 1333,
    size_divisor: int = 32,
    input_data_format: Optional[Union[str, ChannelDimension]] = None,
) -> Tuple[int, int]:
    input_height, input_width = get_image_size(input_image, input_data_format)
    min_size, max_size = shorter, longer

    scale = min_size / min(input_height, input_width)

    if input_height < input_width:
        new_height = min_size
        new_width = scale * input_width
    else:
        new_height = scale * input_height
        new_width = min_size

    if max(new_height, new_width) > max_size:
        scale = max_size / max(new_height, new_width)
        new_height = scale * new_height
        new_width = scale * new_width

    new_height, new_width = int(new_height + 0.5), int(new_width + 0.5)
    new_height = new_height // size_divisor * size_divisor
    new_width = new_width // size_divisor * size_divisor

    return new_height, new_width


class ViltImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        size_divisor: int = 32,
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        do_rescale: bool = True,
        rescale_factor: Union[int, float] = 1 / 255,
        do_normalize: bool = True,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        do_pad: bool = True,
        **kwargs,
    ) -> None:
        if "pad_and_return_pixel_mask" in kwargs:
            do_pad = kwargs.pop("pad_and_return_pixel_mask")

        super().__init__(**kwargs)
        size = size if size is not None else {"shortest_edge": 384}
        size = get_size_dict(size, default_to_square=False)

        self.do_resize = do_resize
        self.size = size
        self.size_divisor = size_divisor
        self.resample = resample
        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_STANDARD_MEAN
        self.image_std = image_std if image_std is not None else IMAGENET_STANDARD_STD
        self.do_pad = do_pad

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

    def resize(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        size_divisor: int = 32,
        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, default_to_square=False)
        if "shortest_edge" not in size:
            raise ValueError(f"The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}")
        shorter = size["shortest_edge"]
        longer = int(1333 / 800 * shorter)
        output_size = get_resize_output_image_size(
            image, shorter=shorter, longer=longer, size_divisor=size_divisor, input_data_format=input_data_format
        )
        return resize(
            image,
            size=output_size,
            resample=resample,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )

    
    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 preprocess(
        self,
        images: ImageInput,
        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,
        do_pad: Optional[bool] = 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_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
        do_pad = do_pad if do_pad is not None else self.do_pad

        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=False)

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

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

        
        images = [to_numpy_array(image) for image in images]

        if is_scaled_image(images[0]) 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(images[0])

        if do_resize:
            images = [
                self.resize(
                    image=image,
                    size=size,
                    size_divisor=size_divisor,
                    resample=resample,
                    input_data_format=input_data_format,
                )
                for image in images
            ]

        if do_rescale:
            images = [
                self.rescale(image=image, scale=rescale_factor, input_data_format=input_data_format)
                for image in images
            ]

        if do_normalize:
            images = [
                self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
                for image in images
            ]

        images = [
            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images
        ]

        if do_pad:
            encoded_outputs = self.pad(
                images, return_pixel_mask=True, return_tensors=return_tensors, input_data_format=data_format
            )
        else:
            encoded_outputs = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)

        return encoded_outputs