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 levit from typing import dict iterable optional union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import getresizeoutputimagesize resize tochanneldimensionformat from imageutils import imagenetdefaultmean imagenetdefaultstd channeldimension imageinput pilimageresampling inferchanneldimensionformat isscaledimage makelistofimages tonumpyarray validimages from utils import tensortype logging logger logging getloggername class levitimageprocessorbaseimageprocessor r constructs a levit image processor args doresize bool optional defaults to true wwhether to resize the shortest edge of the input to int256224 size can be overridden by the doresize parameter in the preprocess method size dictstr int optional defaults to shortestedge 224 size of the output image after resizing if size is a dict with keys width and height the image will be resized to sizeheight sizewidth if size is a dict with key shortestedge the shortest edge value c is rescaled to intc 256224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to sizeshortestegde height width sizeshortestegde can be overridden by the size parameter in the preprocess method resample pilimageresampling optional defaults to resampling bicubic 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 or not to center crop the input to cropsizeheight cropsizewidth can be overridden by the docentercrop parameter in the preprocess method cropsize dict optional defaults to height 224 width 224 desired image size after centercrop can be overridden by the cropsize parameter in the preprocess method dorescale bool optional defaults to true controls 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 scale factor to use if rescaling the image can be overridden by the rescalefactor parameter in the preprocess method donormalize bool optional defaults to true controls whether to normalize the image can be overridden by the donormalize parameter in the preprocess method imagemean listint optional defaults to 0 485 0 456 0 406 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 listint optional defaults to 0 229 0 224 0 225 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 if size is a dict with keys width and height the image will be resized to sizeheight sizewidth if size is a dict with key shortestedge the shortest edge value c is rescaled to intc 256224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to sizeshortestegde height width sizeshortestegde args image np ndarray image to resize size dictstr int size of the output image after resizing if size is a dict with keys width and height the image will be resized to height width if size is a dict with key shortestedge the shortest edge value c is rescaled to intc 256224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to size height width size resample pilimageresampling 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 channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred sizedict is a dict with either keys height and width or shortestedge preprocess an image or batch of images to be used as input to a levit model args images imageinput image or batch of images 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 size of the output image after resizing if size is a dict with keys width and height the image will be resized to height width if size is a dict with key shortestedge the shortest edge value c is rescaled to intc 256224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to size height width size resample pilimageresampling optional defaults to pilimageresampling bicubic resampling filter to use when resiizing the image docentercrop bool optional defaults to self docentercrop whether to center crop the image cropsize dictstr int optional defaults to self cropsize size of the output image after center cropping crops images to cropsizeheight cropsizewidth dorescale bool optional defaults to self dorescale whether to rescale the image pixel values by rescalingfactor typical to values between 0 and 1 rescalefactor float optional defaults to self rescalefactor factor to rescale the image pixel values by donormalize bool optional defaults to self donormalize whether to normalize the image pixel values by imagemean and imagestd imagemean float or listfloat optional defaults to self imagemean mean to normalize the image pixel values by imagestd float or listfloat optional defaults to self imagestd standard deviation to normalize the image pixel values by 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 defaults to channeldimension first the channel dimension format for the output image if unset the channel dimension format of the input image is used can be one of channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format inputdataformat 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 levit constructs a levit image processor args do_resize bool optional defaults to true wwhether to resize the shortest edge of the input to int 256 224 size can be overridden by the do_resize parameter in the preprocess method size dict str int optional defaults to shortest_edge 224 size of the output image after resizing if size is a dict with keys width and height the image will be resized to size height size width if size is a dict with key shortest_edge the shortest edge value c is rescaled to int c 256 224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to size shortest_egde height width size shortest_egde can be overridden by the size parameter in the preprocess method resample pilimageresampling optional defaults to resampling bicubic 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 or not to center crop the input to crop_size height crop_size width can be overridden by the do_center_crop parameter in the preprocess method crop_size dict optional defaults to height 224 width 224 desired image size after center_crop can be overridden by the crop_size parameter in the preprocess method do_rescale bool optional defaults to true controls 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 scale factor to use if rescaling the image can be overridden by the rescale_factor parameter in the preprocess method do_normalize bool optional defaults to true controls whether to normalize the image can be overridden by the do_normalize parameter in the preprocess method image_mean list int optional defaults to 0 485 0 456 0 406 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 list int optional defaults to 0 229 0 224 0 225 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 if size is a dict with keys width and height the image will be resized to size height size width if size is a dict with key shortest_edge the shortest edge value c is rescaled to int c 256 224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to size shortest_egde height width size shortest_egde args image np ndarray image to resize size dict str int size of the output image after resizing if size is a dict with keys width and height the image will be resized to height width if size is a dict with key shortest_edge the shortest edge value c is rescaled to int c 256 224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to size height width size resample pilimageresampling 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 channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred size_dict is a dict with either keys height and width or shortest_edge preprocess an image or batch of images to be used as input to a levit model args images imageinput image or batch of images 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 output image after resizing if size is a dict with keys width and height the image will be resized to height width if size is a dict with key shortest_edge the shortest edge value c is rescaled to int c 256 224 the smaller edge of the image will be matched to this value i e if height width then image will be rescaled to size height width size resample pilimageresampling optional defaults to pilimageresampling bicubic resampling filter to use when resiizing the image 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 output image after center cropping crops images to crop_size height crop_size width do_rescale bool optional defaults to self do_rescale whether to rescale the image pixel values by rescaling_factor typical to values between 0 and 1 rescale_factor float optional defaults to self rescale_factor factor to rescale the image pixel values by do_normalize bool optional defaults to self do_normalize whether to normalize the image pixel values by image_mean and image_std image_mean float or list float optional defaults to self image_mean mean to normalize the image pixel values by image_std float or list float optional defaults to self image_std standard deviation to normalize the image pixel values by 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 defaults to channeldimension first the channel dimension format for the output image if unset the channel dimension format of the input image is used can be one of channels_first or channeldimension first image in num_channels height width format channels_last or channeldimension last image in height width num_channels format input_data_format 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 Dict, Iterable, Optional, Union

import numpy as np

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


logger = logging.get_logger(__name__)


class LevitImageProcessor(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]]] = IMAGENET_DEFAULT_MEAN,
        image_std: Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_STD,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        size = size if size is not None else {"shortest_edge": 224}
        size = get_size_dict(size, default_to_square=False)
        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")

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

    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_dict = get_size_dict(size, default_to_square=False)
        
        if "shortest_edge" in size:
            shortest_edge = int((256 / 224) * size["shortest_edge"])
            output_size = get_resize_output_image_size(
                image, size=shortest_edge, default_to_square=False, input_data_format=input_data_format
            )
            size_dict = {"height": output_size[0], "width": output_size[1]}
        if "height" not in size_dict or "width" not in size_dict:
            raise ValueError(
                f"Size dict must have keys 'height' and 'width' or 'shortest_edge'. Got {size_dict.keys()}"
            )
        return resize(
            image,
            size=(size_dict["height"], size_dict["width"]),
            resample=resample,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )

    def preprocess(
        self,
        images: ImageInput,
        do_resize: Optional[bool] = None,
        size: Optional[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, Iterable[float]]] = None,
        image_std: Optional[Union[float, Iterable[float]]] = None,
        return_tensors: Optional[TensorType] = None,
        data_format: ChannelDimension = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> BatchFeature:
        
        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_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

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

        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:
            raise ValueError("Size 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_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, size, resample, input_data_format=input_data_format) for image in images]

        if do_center_crop:
            images = [self.center_crop(image, crop_size, input_data_format=input_data_format) for image in images]

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

        if do_normalize:
            images = [
                self.normalize(image, image_mean, 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
        ]

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