codingutf8 2023 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 nougat from typing import dict list optional union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import getresizeoutputimagesize pad resize tochanneldimensionformat topilimage from imageutils import imagenetdefaultmean imagenetdefaultstd channeldimension imageinput pilimageresampling getimagesize inferchanneldimensionformat isscaledimage makelistofimages tonumpyarray validimages from utils import tensortype logging from utils importutils import iscv2available isvisionavailable logger logging getloggername if iscv2available pass if isvisionavailable import pil class nougatimageprocessorbaseimageprocessor r constructs a nougat image processor args docropmargin bool optional defaults to true whether to crop the image margins doresize bool optional defaults to true whether to resize the image s height width dimensions to the specified size can be overridden by doresize in the preprocess method size dictstr int optional defaults to height 896 width 672 size of the image after resizing can be overridden 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 resample in the preprocess method dothumbnail bool optional defaults to true whether to resize the image using thumbnail method doalignlongaxis bool optional defaults to false whether to align the long axis of the image with the long axis of size by rotating by 90 degrees dopad bool optional defaults to true whether to pad the images to the largest image size in the batch 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 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 whether to normalize the image can be overridden by donormalize in the preprocess method imagemean float or listfloat optional defaults to imagenetdefaultmean 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 imagenetdefaultstd image standard deviation this is a reimplementation of a findnonzero function equivalent to cv2 nonzeroindices np columnstacknp nonzeroimage idxvec nonzeroindices 1 0 idxvec idxvec reshape1 1 2 return idxvec def pythonboundingrectself coordinates crops the margin of the image gray pixels are considered margin i e pixels with a value below the threshold args image np array the image to be cropped graythreshold int optional defaults to 200 value below which pixels are considered to be gray dataformat channeldimension optional the channel dimension format of the output image if unset will use the inferred format from the input inputdataformat channeldimension optional the channel dimension format of the input image if unset will use the inferred format from the input copied from transformers models donut imageprocessingdonut donutimageprocessor alignlongaxis align the long axis of the image to the longest axis of the specified size args image np ndarray the image to be aligned size dictstr int the size height h width w to align the long axis to dataformat str or channeldimension optional the data format of the output image if unset the same format as the input image is used inputdataformat channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred returns np ndarray the aligned image pad the image to the specified size at the top bottom left and right args image np ndarray the image to be padded size dictstr int the size height h width w to pad the image to dataformat str or channeldimension optional the data format of the output image if unset the same format as the input image is used inputdataformat channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred copied from transformers models donut imageprocessingdonut donutimageprocessor thumbnail resize the image to make a thumbnail the image is resized so that no dimension is larger than any corresponding dimension of the specified size args image np ndarray the image to be resized size dictstr int the size height h width w to resize the image to resample pilimageresampling optional defaults to pilimageresampling bicubic the resampling filter to use dataformat optionalunionstr channeldimension optional the data format of the output image if unset the same format as the input image is used inputdataformat channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred we always resize to the smallest of either the input or output size copied from transformers models donut imageprocessingdonut donutimageprocessor resize resizes image to height width specified by size using the pil library args image np ndarray image to resize size dictstr int size of the output image 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 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 docropmargin bool optional defaults to self docropmargin whether to crop the image margins 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 shortest edge of the image is resized to minsizeheight sizewidth with the longest edge resized to keep the input aspect ratio 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 dothumbnail bool optional defaults to self dothumbnail whether to resize the image using thumbnail method doalignlongaxis bool optional defaults to self doalignlongaxis whether to align the long axis of the image with the long axis of size by rotating by 90 degrees dopad bool optional defaults to self dopad whether to pad the images to the largest image size in the batch dorescale bool optional defaults to self dorescale whether to rescale the image by the specified scale rescalefactor rescalefactor int or float optional defaults to self rescalefactor scale factor to use if rescaling the image donormalize bool optional defaults to self donormalize whether to normalize the image imagemean float or listfloat optional defaults to self imagemean image mean to use for normalization imagestd float or listfloat optional defaults to self imagestd image standard deviation to use for normalization 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 unset defaults to the channel dimension format of the input image 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 2023 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 nougat constructs a nougat image processor args do_crop_margin bool optional defaults to true whether to crop the image margins do_resize bool optional defaults to true whether to resize the image s height width dimensions to the specified size can be overridden by do_resize in the preprocess method size dict str int optional defaults to height 896 width 672 size of the image after resizing can be overridden 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 resample in the preprocess method do_thumbnail bool optional defaults to true whether to resize the image using thumbnail method do_align_long_axis bool optional defaults to false whether to align the long axis of the image with the long axis of size by rotating by 90 degrees do_pad bool optional defaults to true whether to pad the images to the largest image size in the batch 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 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 whether to normalize the image can be overridden by do_normalize in the preprocess method image_mean float or list float optional defaults to imagenet_default_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_default_std image standard deviation this is a reimplementation of a findnonzero function equivalent to cv2 this is a reimplementation of a boundingrect function equivalent to cv2 crops the margin of the image gray pixels are considered margin i e pixels with a value below the threshold args image np array the image to be cropped gray_threshold int optional defaults to 200 value below which pixels are considered to be gray data_format channeldimension optional the channel dimension format of the output image if unset will use the inferred format from the input input_data_format channeldimension optional the channel dimension format of the input image if unset will use the inferred format from the input copied from transformers models donut image_processing_donut donutimageprocessor align_long_axis align the long axis of the image to the longest axis of the specified size args image np ndarray the image to be aligned size dict str int the size height h width w to align the long axis to data_format str or channeldimension optional the data format of the output image if unset the same format as the input image is used input_data_format channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred returns np ndarray the aligned image pad the image to the specified size at the top bottom left and right args image np ndarray the image to be padded size dict str int the size height h width w to pad the image to data_format str or channeldimension optional the data format of the output image if unset the same format as the input image is used input_data_format channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred copied from transformers models donut image_processing_donut donutimageprocessor thumbnail resize the image to make a thumbnail the image is resized so that no dimension is larger than any corresponding dimension of the specified size args image np ndarray the image to be resized size dict str int the size height h width w to resize the image to resample pilimageresampling optional defaults to pilimageresampling bicubic the resampling filter to use data_format optional union str channeldimension optional the data format of the output image if unset the same format as the input image is used input_data_format channeldimension or str optional the channel dimension format of the input image if not provided it will be inferred we always resize to the smallest of either the input or output size copied from transformers models donut image_processing_donut donutimageprocessor resize resizes image to height width specified by size using the pil library args image np ndarray image to resize size dict str int size of the output image 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 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 do_crop_margin bool optional defaults to self do_crop_margin whether to crop the image margins 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 shortest edge of the image is resized to min size height size width with the longest edge resized to keep the input aspect ratio 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_thumbnail bool optional defaults to self do_thumbnail whether to resize the image using thumbnail method do_align_long_axis bool optional defaults to self do_align_long_axis whether to align the long axis of the image with the long axis of size by rotating by 90 degrees do_pad bool optional defaults to self do_pad whether to pad the images to the largest image size in the batch do_rescale bool optional defaults to self do_rescale whether to rescale the image by the specified scale rescale_factor rescale_factor int or float optional defaults to self rescale_factor scale factor to use if rescaling the image 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 use for normalization image_std float or list float optional defaults to self image_std image standard deviation to use for normalization 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 defaults to the 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 all transformations expect numpy arrays we assume that all images have the same channel dimension format
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 (
    get_resize_output_image_size,
    pad,
    resize,
    to_channel_dimension_format,
    to_pil_image,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_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, logging
from ...utils.import_utils import is_cv2_available, is_vision_available


logger = logging.get_logger(__name__)


if is_cv2_available():
    pass


if is_vision_available():
    import PIL


class NougatImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_crop_margin: bool = True,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        do_thumbnail: bool = True,
        do_align_long_axis: bool = False,
        do_pad: bool = True,
        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,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        size = size if size is not None else {"height": 896, "width": 672}
        size = get_size_dict(size)

        self.do_crop_margin = do_crop_margin
        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.do_thumbnail = do_thumbnail
        self.do_align_long_axis = do_align_long_axis
        self.do_pad = do_pad
        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 python_find_non_zero(self, image: np.array):
        
        non_zero_indices = np.column_stack(np.nonzero(image))
        idxvec = non_zero_indices[:, [1, 0]]
        idxvec = idxvec.reshape(-1, 1, 2)
        return idxvec

    def python_bounding_rect(self, coordinates):
        
        min_values = np.min(coordinates, axis=(0, 1)).astype(int)
        max_values = np.max(coordinates, axis=(0, 1)).astype(int)
        x_min, y_min = min_values[0], min_values[1]
        width = max_values[0] - x_min + 1
        height = max_values[1] - y_min + 1
        return x_min, y_min, width, height

    def crop_margin(
        self,
        image: np.array,
        gray_threshold: int = 200,
        data_format: Optional[ChannelDimension] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.array:
        
        if input_data_format is None:
            input_data_format = infer_channel_dimension_format(image)

        image = to_pil_image(image, input_data_format=input_data_format)
        data = np.array(image.convert("L")).astype(np.uint8)
        max_val = data.max()
        min_val = data.min()
        if max_val == min_val:
            image = np.array(image)
            image = (
                to_channel_dimension_format(image, data_format, input_data_format)
                if data_format is not None
                else image
            )
            return image
        data = (data - min_val) / (max_val - min_val) * 255
        gray = data < gray_threshold
        coords = self.python_find_non_zero(gray)
        x_min, y_min, width, height = self.python_bounding_rect(coords)
        image = image.crop((x_min, y_min, x_min + width, y_min + height))
        image = np.array(image).astype(np.uint8)
        image = to_channel_dimension_format(image, input_data_format, ChannelDimension.LAST)

        image = (
            to_channel_dimension_format(image, data_format, input_data_format) if data_format is not None else image
        )

        return image

    
    def align_long_axis(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        data_format: Optional[Union[str, 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 = size["height"], size["width"]

        if (output_width < output_height and input_width > input_height) or (
            output_width > output_height and input_width < input_height
        ):
            image = np.rot90(image, 3)

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

        return image

    def pad_image(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
    ) -> np.ndarray:
        
        output_height, output_width = size["height"], size["width"]
        input_height, input_width = get_image_size(image, channel_dim=input_data_format)

        delta_width = output_width - input_width
        delta_height = output_height - input_height

        pad_top = delta_height // 2
        pad_left = delta_width // 2

        pad_bottom = delta_height - pad_top
        pad_right = delta_width - pad_left

        padding = ((pad_top, pad_bottom), (pad_left, pad_right))
        return pad(image, padding, data_format=data_format, input_data_format=input_data_format)

    
    def thumbnail(
        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:
        
        input_height, input_width = get_image_size(image, channel_dim=input_data_format)
        output_height, output_width = size["height"], size["width"]

        
        height = min(input_height, output_height)
        width = min(input_width, output_width)

        if height == input_height and width == input_width:
            return image

        if input_height > input_width:
            width = int(input_width * height / input_height)
        elif input_width > input_height:
            height = int(input_height * width / input_width)

        return resize(
            image,
            size=(height, width),
            resample=resample,
            reducing_gap=2.0,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )

    
    def resize(
        self,
        image: np.ndarray,
        size: Dict[str, int],
        resample: PILImageResampling = PILImageResampling.BICUBIC,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> np.ndarray:
        
        size = get_size_dict(size)
        shortest_edge = min(size["height"], size["width"])
        output_size = get_resize_output_image_size(
            image, size=shortest_edge, default_to_square=False, input_data_format=input_data_format
        )
        resized_image = resize(
            image,
            size=output_size,
            resample=resample,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )
        return resized_image

    def preprocess(
        self,
        images: ImageInput,
        do_crop_margin: bool = None,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        do_thumbnail: bool = None,
        do_align_long_axis: bool = None,
        do_pad: bool = None,
        do_rescale: bool = None,
        rescale_factor: Union[int, float] = None,
        do_normalize: 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: Optional[ChannelDimension] = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> PIL.Image.Image:
        
        do_crop_margin = do_crop_margin if do_crop_margin is not None else self.do_crop_margin
        do_resize = do_resize if do_resize is not None else self.do_resize
        size = size if size is not None else self.size
        resample = resample if resample is not None else self.resample
        do_thumbnail = do_thumbnail if do_thumbnail is not None else self.do_thumbnail
        do_align_long_axis = do_align_long_axis if do_align_long_axis is not None else self.do_align_long_axis
        do_pad = do_pad if do_pad is not None else self.do_pad
        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

        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_pad and size is None:
            raise ValueError("Size must be specified if do_pad 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_crop_margin:
            images = [self.crop_margin(image, input_data_format=input_data_format) for image in images]

        if do_align_long_axis:
            images = [self.align_long_axis(image, size=size, input_data_format=input_data_format) for image in images]

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

        if do_thumbnail:
            images = [self.thumbnail(image=image, size=size, input_data_format=input_data_format) for image in images]

        if do_pad:
            images = [self.pad_image(image=image, size=size, 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
        ]

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