codingutf8 2023 the intel aia team s and 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 tvp from typing import dict iterable list optional tuple union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import paddingmode flipchannelorder pad resize tochanneldimensionformat from imageutils import imagenetstandardmean imagenetstandardstd channeldimension imageinput pilimageresampling getimagesize isvalidimage tonumpyarray validimages from utils import tensortype isvisionavailable logging if isvisionavailable import pil logger logging getloggername copied from transformers models vivit imageprocessingvivit makebatched def makebatchedvideos listlistimageinput if isinstancevideos list tuple and isinstancevideos0 list tuple and isvalidimagevideos00 return videos elif isinstancevideos list tuple and isvalidimagevideos0 return videos elif isvalidimagevideos return videos raise valueerrorfcould not make batched video from videos def getresizeoutputimagesize inputimage np ndarray maxsize int 448 inputdataformat optionalunionstr channeldimension none tupleint int height width getimagesizeinputimage inputdataformat if height width ratio width 1 0 height newheight maxsize newwidth newheight ratio else ratio height 1 0 width newwidth maxsize newheight newwidth ratio size intnewheight intnewwidth return size class tvpimageprocessorbaseimageprocessor r constructs a tvp 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 longestedge 448 size of the output image after resizing the longest edge of the image will be resized to sizelongestedge while maintaining the aspect ratio of the original image can be overriden by size in the preprocess method resample pilimageresampling optional defaults to resampling bilinear resampling filter to use if resizing the image can be overridden by the resample parameter in the preprocess method docentercrop bool optional defaults to true whether to center crop the image to the specified cropsize can be overridden by the docentercrop parameter in the preprocess method cropsize dictstr int optional defaults to height 448 width 448 size of the image after applying the center crop can be overridden by the cropsize parameter in the preprocess method dorescale bool optional defaults to true whether 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 defines the scale factor to use if rescaling the image can be overridden by the rescalefactor parameter in the preprocess method dopad bool optional defaults to true whether to pad the image can be overridden by the dopad parameter in the preprocess method padsize dictstr int optional defaults to height 448 width 448 size of the image after applying the padding can be overridden by the padsize parameter in the preprocess method constantvalues unionfloat iterablefloat optional defaults to 0 the fill value to use when padding the image padmode paddingmode optional defaults to paddingmode constant use what kind of mode in padding donormalize bool optional defaults to true whether to normalize the image can be overridden by the donormalize parameter in the preprocess method doflipchannelorder bool optional defaults to true whether to flip the color channels from rgb to bgr can be overridden by the doflipchannelorder 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 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 resize an image args image np ndarray image to resize size dictstr int size of the output image if size is of the form height h width w the output image will have the size h w if size is of the form longestedge s the output image will have its longest edge of length s while keeping the aspect ratio of the original image resample pilimageresampling optional defaults to pilimageresampling bilinear 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 pad an image with zeros to the given size args image np ndarray image to pad padsize dictstr int size of the output image with pad constantvalues unionfloat iterablefloat the fill value to use when padding the image padmode paddingmode the pad mode default to paddingmode constant dataformat channeldimension or str 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 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 docentercrop and cropsize is none raise valueerrorcrop size must be specified if docentercrop is true if dorescale and rescalefactor is none raise valueerrorrescale factor must be specified if dorescale is true if dopad and padsize is none raise valueerrorpadding size must be specified if dopad 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 doresize image self resizeimageimage sizesize resampleresample inputdataformatinputdataformat if docentercrop image self centercropimage sizecropsize inputdataformatinputdataformat if dorescale image self rescaleimageimage scalerescalefactor inputdataformatinputdataformat if donormalize image self normalize imageimage astypenp float32 meanimagemean stdimagestd inputdataformatinputdataformat if dopad image self padimage imageimage padsizepadsize constantvaluesconstantvalues padmodepadmode inputdataformatinputdataformat the pretrained checkpoints assume images are bgr not rgb if doflipchannelorder image flipchannelorderimageimage inputdataformatinputdataformat image tochanneldimensionformatimage dataformat inputchanneldiminputdataformat return image def preprocess self videos unionimageinput listimageinput listlistimageinput doresize bool none size dictstr int none resample pilimageresampling none docentercrop bool none cropsize dictstr int none dorescale bool none rescalefactor float none dopad bool none padsize dictstr int none constantvalues unionfloat iterablefloat none padmode paddingmode none donormalize bool none doflipchannelorder bool none imagemean optionalunionfloat listfloat none imagestd optionalunionfloat listfloat 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 resample resample if resample is not none else self resample docentercrop docentercrop if docentercrop is not none else self docentercrop dorescale dorescale if dorescale is not none else self dorescale rescalefactor rescalefactor if rescalefactor is not none else self rescalefactor dopad dopad if dopad is not none else self dopad padsize padsize if padsize is not none else self padsize constantvalues constantvalues if constantvalues is not none else self constantvalues padmode padmode if padmode else self padmode donormalize donormalize if donormalize is not none else self donormalize doflipchannelorder doflipchannelorder if doflipchannelorder is not none else self doflipchannelorder imagemean imagemean if imagemean is not none else self imagemean imagestd imagestd if imagestd is not none else self imagestd size size if size is not none else self size size getsizedictsize defaulttosquarefalse cropsize cropsize if cropsize is not none else self cropsize cropsize getsizedictcropsize paramnamecropsize if not validimagesvideos raise valueerror invalid image type must be of type pil image image numpy ndarray torch tensor tf tensor or jax ndarray videos makebatchedvideos videos np array self preprocessimage imageimg doresizedoresize sizesize resampleresample docentercropdocentercrop cropsizecropsize dorescaledorescale rescalefactorrescalefactor dopaddopad padsizepadsize constantvaluesconstantvalues padmodepadmode donormalizedonormalize doflipchannelorderdoflipchannelorder imagemeanimagemean imagestdimagestd dataformatdataformat inputdataformatinputdataformat for img in video for video in videos data pixelvalues videos return batchfeaturedatadata tensortypereturntensors coding utf 8 2023 the intel aia team s and 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 tvp copied from transformers models vivit image_processing_vivit make_batched constructs a tvp 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 longest_edge 448 size of the output image after resizing the longest edge of the image will be resized to size longest_edge while maintaining the aspect ratio of the original image can be overriden by size in the preprocess method resample pilimageresampling optional defaults to resampling bilinear resampling filter to use if resizing the image can be overridden by the resample parameter in the preprocess method do_center_crop bool optional defaults to true whether to center crop the image to the specified crop_size can be overridden by the do_center_crop parameter in the preprocess method crop_size dict str int optional defaults to height 448 width 448 size of the image after applying the center crop can be overridden by the crop_size parameter in the preprocess method 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 the preprocess method rescale_factor int or float optional defaults to 1 255 defines the scale factor to use if rescaling the image can be overridden by the rescale_factor parameter in the preprocess method do_pad bool optional defaults to true whether to pad the image can be overridden by the do_pad parameter in the preprocess method pad_size dict str int optional defaults to height 448 width 448 size of the image after applying the padding can be overridden by the pad_size parameter in the preprocess method constant_values union float iterable float optional defaults to 0 the fill value to use when padding the image pad_mode paddingmode optional defaults to paddingmode constant use what kind of mode in padding do_normalize bool optional defaults to true whether to normalize the image can be overridden by the do_normalize parameter in the preprocess method do_flip_channel_order bool optional defaults to true whether to flip the color channels from rgb to bgr can be overridden by the do_flip_channel_order 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 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 resize an image args image np ndarray image to resize size dict str int size of the output image if size is of the form height h width w the output image will have the size h w if size is of the form longest_edge s the output image will have its longest edge of length s while keeping the aspect ratio of the original image resample pilimageresampling optional defaults to pilimageresampling bilinear 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 pad an image with zeros to the given size args image np ndarray image to pad pad_size dict str int size of the output image with pad constant_values union float iterable float the fill value to use when padding the image pad_mode paddingmode the pad mode default to paddingmode constant data_format channeldimension or str 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 preprocesses a single image all transformations expect numpy arrays the pretrained checkpoints assume images are bgr not rgb preprocess an image or batch of images args videos imageinput or list imageinput or list list imageinput frames to preprocess 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 after applying resize resample pilimageresampling 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_centre_crop whether to centre crop the image crop_size dict str int optional defaults to self crop_size size of the image after applying the centre crop 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_pad bool optional defaults to true whether to pad the image can be overridden by the do_pad parameter in the preprocess method pad_size dict str int optional defaults to height 448 width 448 size of the image after applying the padding can be overridden by the pad_size parameter in the preprocess method constant_values union float iterable float optional defaults to 0 the fill value to use when padding the image pad_mode paddingmode optional defaults to paddingmode constant use what kind of mode in padding do_normalize bool optional defaults to self do_normalize whether to normalize the image do_flip_channel_order bool optional defaults to self do_flip_channel_order whether to flip the channel order of 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_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 unset use the inferred channel dimension format of the input image 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
from typing import 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,
    flip_channel_order,
    pad,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    get_image_size,
    is_valid_image,
    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 make_batched(videos) -> List[List[ImageInput]]:
    if isinstance(videos, (list, tuple)) and isinstance(videos[0], (list, tuple)) and is_valid_image(videos[0][0]):
        return videos

    elif isinstance(videos, (list, tuple)) and is_valid_image(videos[0]):
        return [videos]

    elif is_valid_image(videos):
        return [[videos]]

    raise ValueError(f"Could not make batched video from {videos}")


def get_resize_output_image_size(
    input_image: np.ndarray,
    max_size: int = 448,
    input_data_format: Optional[Union[str, ChannelDimension]] = None,
) -> Tuple[int, int]:
    height, width = get_image_size(input_image, input_data_format)
    if height >= width:
        ratio = width * 1.0 / height
        new_height = max_size
        new_width = new_height * ratio
    else:
        ratio = height * 1.0 / width
        new_width = max_size
        new_height = new_width * ratio
    size = (int(new_height), int(new_width))

    return size


class TvpImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        do_center_crop: bool = True,
        crop_size: Dict[str, int] = None,
        do_rescale: bool = True,
        rescale_factor: Union[int, float] = 1 / 255,
        do_pad: bool = True,
        pad_size: Dict[str, int] = None,
        constant_values: Union[float, Iterable[float]] = 0,
        pad_mode: PaddingMode = PaddingMode.CONSTANT,
        do_normalize: bool = True,
        do_flip_channel_order: bool = True,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        size = size if size is not None else {"longest_edge": 448}
        crop_size = crop_size if crop_size is not None else {"height": 448, "width": 448}
        pad_size = pad_size if pad_size is not None else {"height": 448, "width": 448}

        self.do_resize = do_resize
        self.size = size
        self.do_center_crop = do_center_crop
        self.crop_size = crop_size
        self.resample = resample
        self.do_rescale = do_rescale
        self.rescale_factor = rescale_factor
        self.do_pad = do_pad
        self.pad_size = pad_size
        self.constant_values = constant_values
        self.pad_mode = pad_mode
        self.do_normalize = do_normalize
        self.do_flip_channel_order = do_flip_channel_order
        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

    def resize(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        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 "height" in size and "width" in size:
            output_size = (size["height"], size["width"])
        elif "longest_edge" in size:
            output_size = get_resize_output_image_size(image, size["longest_edge"], input_data_format)
        else:
            raise ValueError(f"Size must have 'height' and 'width' or 'longest_edge' as keys. Got {size.keys()}")

        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,
        pad_size: Dict[str, int] = None,
        constant_values: Union[float, Iterable[float]] = 0,
        pad_mode: PaddingMode = PaddingMode.CONSTANT,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ):
        
        height, width = get_image_size(image, channel_dim=input_data_format)
        max_height = pad_size.get("height", height)
        max_width = pad_size.get("width", width)

        pad_right, pad_bottom = max_width - width, max_height - height
        if pad_right < 0 or pad_bottom < 0:
            raise ValueError("The padding size must be greater than image size")

        padding = ((0, pad_bottom), (0, pad_right))
        padded_image = pad(
            image,
            padding,
            mode=pad_mode,
            constant_values=constant_values,
            data_format=data_format,
            input_data_format=input_data_format,
        )

        return padded_image

    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_pad: bool = True,
        pad_size: Dict[str, int] = None,
        constant_values: Union[float, Iterable[float]] = None,
        pad_mode: PaddingMode = None,
        do_normalize: bool = None,
        do_flip_channel_order: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        data_format: Optional[ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> 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_center_crop and crop_size is None:
            raise ValueError("Crop size must be specified if do_center_crop is True.")

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

        if do_pad and pad_size is None:
            raise ValueError("Padding size must be specified if do_pad 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 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, 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.astype(np.float32), mean=image_mean, std=image_std, input_data_format=input_data_format
            )

        if do_pad:
            image = self.pad_image(
                image=image,
                pad_size=pad_size,
                constant_values=constant_values,
                pad_mode=pad_mode,
                input_data_format=input_data_format,
            )

        
        if do_flip_channel_order:
            image = flip_channel_order(image=image, input_data_format=input_data_format)

        image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)

        return image

    def preprocess(
        self,
        videos: Union[ImageInput, List[ImageInput], List[List[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_pad: bool = None,
        pad_size: Dict[str, int] = None,
        constant_values: Union[float, Iterable[float]] = None,
        pad_mode: PaddingMode = None,
        do_normalize: bool = None,
        do_flip_channel_order: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[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
        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
        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_pad = do_pad if do_pad is not None else self.do_pad
        pad_size = pad_size if pad_size is not None else self.pad_size
        constant_values = constant_values if constant_values is not None else self.constant_values
        pad_mode = pad_mode if pad_mode else self.pad_mode
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        do_flip_channel_order = (
            do_flip_channel_order if do_flip_channel_order is not None else self.do_flip_channel_order
        )
        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

        size = size if size is not None else self.size
        size = get_size_dict(size, default_to_square=False)
        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")

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

        videos = make_batched(videos)

        videos = [
            np.array(
                [
                    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_pad=do_pad,
                        pad_size=pad_size,
                        constant_values=constant_values,
                        pad_mode=pad_mode,
                        do_normalize=do_normalize,
                        do_flip_channel_order=do_flip_channel_order,
                        image_mean=image_mean,
                        image_std=image_std,
                        data_format=data_format,
                        input_data_format=input_data_format,
                    )
                    for img in video
                ]
            )
            for video in videos
        ]

        data = {"pixel_values": videos}
        return BatchFeature(data=data, tensor_type=return_tensors)