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 imagegpt from typing import dict list optional union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import rescale resize tochanneldimensionformat from imageutils import channeldimension imageinput pilimageresampling inferchanneldimensionformat isscaledimage makelistofimages tonumpyarray validimages from utils import tensortype isvisionavailable logging if isvisionavailable import pil logger logging getloggername def squaredeuclideandistancea b b b t a2 np sumnp squarea axis1 b2 np sumnp squareb axis0 ab np matmula b d a2 none 2 ab b2none return d def colorquantizex clusters x x reshape1 3 d squaredeuclideandistancex clusters return np argmind axis1 class imagegptimageprocessorbaseimageprocessor r constructs a imagegpt image processor this image processor can be used to resize images to a smaller resolution such as 32x32 or 64x64 normalize them and finally color quantize them to obtain sequences of pixel values color clusters args clusters np ndarray or listlistint optional the color clusters to use of shape nclusters 3 when color quantizing can be overriden by clusters in preprocess doresize bool optional defaults to true whether to resize the image s dimensions to sizeheight sizewidth can be overridden by doresize in preprocess size dictstr int optional defaults to height 256 width 256 size of the image after resizing can be overridden by size in preprocess resample pilimageresampling optional defaults to resampling bilinear resampling filter to use if resizing the image can be overridden by resample in preprocess donormalize bool optional defaults to true whether to normalize the image pixel value to between 1 1 can be overridden by donormalize in preprocess docolorquantize bool optional defaults to true whether to color quantize the image can be overridden by docolorquantize in preprocess clusters is a first argument to maintain backwards compatibility with the old imagegptimageprocessor copied from transformers models vit imageprocessingvit vitimageprocessor resize 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 bilinear pilimageresampling filter to use when resizing the image e g pilimageresampling bilinear 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 normalizes an images pixel values to between 1 1 args image np ndarray image to normalize 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 donormalizefalse doresize bool optional defaults to self doresize whether to resize the image size dictstr int optional defaults to self size size of the image after resizing 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 doresize is set to true donormalize bool optional defaults to self donormalize whether to normalize the image docolorquantize bool optional defaults to self docolorquantize whether to color quantize the image clusters np ndarray or listlistint optional defaults to self clusters clusters used to quantize the image of shape nclusters 3 only has an effect if docolorquantize is set to true 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 only has an effect if docolorquantize is set to false 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 color quantize from batchsize height width 3 to batchsize height width flatten to batchsize heightwidth we need to convert back to a list of images to keep consistent behaviour across processors 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 imagegpt constructs a imagegpt image processor this image processor can be used to resize images to a smaller resolution such as 32x32 or 64x64 normalize them and finally color quantize them to obtain sequences of pixel values color clusters args clusters np ndarray or list list int optional the color clusters to use of shape n_clusters 3 when color quantizing can be overriden by clusters in preprocess do_resize bool optional defaults to true whether to resize the image s dimensions to size height size width can be overridden by do_resize in preprocess size dict str int optional defaults to height 256 width 256 size of the image after resizing can be overridden by size in preprocess resample pilimageresampling optional defaults to resampling bilinear resampling filter to use if resizing the image can be overridden by resample in preprocess do_normalize bool optional defaults to true whether to normalize the image pixel value to between 1 1 can be overridden by do_normalize in preprocess do_color_quantize bool optional defaults to true whether to color quantize the image can be overridden by do_color_quantize in preprocess clusters is a first argument to maintain backwards compatibility with the old imagegptimageprocessor copied from transformers models vit image_processing_vit vitimageprocessor resize 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 bilinear pilimageresampling filter to use when resizing the image e g pilimageresampling bilinear 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 normalizes an images pixel values to between 1 1 args image np ndarray image to normalize 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_normalize 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 after resizing 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_normalize bool optional defaults to self do_normalize whether to normalize the image do_color_quantize bool optional defaults to self do_color_quantize whether to color quantize the image clusters np ndarray or list list int optional defaults to self clusters clusters used to quantize the image of shape n_clusters 3 only has an effect if do_color_quantize 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 only has an effect if do_color_quantize is set to false 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 color quantize from batch_size height width 3 to batch_size height width flatten to batch_size height width we need to convert back to a list of images to keep consistent behaviour across processors
from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    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__)


def squared_euclidean_distance(a, b):
    b = b.T
    a2 = np.sum(np.square(a), axis=1)
    b2 = np.sum(np.square(b), axis=0)
    ab = np.matmul(a, b)
    d = a2[:, None] - 2 * ab + b2[None, :]
    return d


def color_quantize(x, clusters):
    x = x.reshape(-1, 3)
    d = squared_euclidean_distance(x, clusters)
    return np.argmin(d, axis=1)


class ImageGPTImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        
        clusters: Optional[Union[List[List[int]], np.ndarray]] = None,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        do_normalize: bool = True,
        do_color_quantize: bool = True,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        size = size if size is not None else {"height": 256, "width": 256}
        size = get_size_dict(size)
        self.clusters = np.array(clusters) if clusters is not None else None
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_normalize = do_normalize
        self.do_color_quantize = do_color_quantize

    
    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)
        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 normalize(
        self,
        image: np.ndarray,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        image = rescale(image=image, scale=1 / 127.5, data_format=data_format, input_data_format=input_data_format)
        image = image - 1
        return image

    def preprocess(
        self,
        images: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_normalize: bool = None,
        do_color_quantize: Optional[bool] = None,
        clusters: Optional[Union[List[List[int]], np.ndarray]] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: Optional[Union[str, 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_normalize = do_normalize if do_normalize is not None else self.do_normalize
        do_color_quantize = do_color_quantize if do_color_quantize is not None else self.do_color_quantize
        clusters = clusters if clusters is not None else self.clusters
        clusters = np.array(clusters)

        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_color_quantize and clusters is None:
            raise ValueError("Clusters must be specified if do_color_quantize is True.")

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

        if is_scaled_image(images[0]) and do_normalize:
            logger.warning_once(
                "It looks like you are trying to rescale already rescaled images. If you wish to do this, "
                "make sure to set `do_normalize` to `False` and that pixel values are between [-1, 1].",
            )

        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, resample=resample, input_data_format=input_data_format)
                for image in images
            ]

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

        if do_color_quantize:
            images = [to_channel_dimension_format(image, ChannelDimension.LAST, input_data_format) for image in images]
            
            images = np.array(images)
            images = color_quantize(images, clusters).reshape(images.shape[:-1])

            
            batch_size = images.shape[0]
            images = images.reshape(batch_size, -1)

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

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