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 pix2struct import io import math from typing import dict optional union import numpy as np from huggingfacehub import hfhubdownload from imageprocessingutils import baseimageprocessor batchfeature from imagetransforms import converttorgb normalize tochanneldimensionformat topilimage from imageutils import channeldimension imageinput getimagesize inferchanneldimensionformat makelistofimages tonumpyarray validimages from utils import tensortype istorchavailable isvisionavailable logging from utils importutils import requiresbackends if isvisionavailable import textwrap from pil import image imagedraw imagefont if istorchavailable import torch from transformers pytorchutils import istorchgreaterorequalthan111 else istorchgreaterorequalthan111 false logger logging getloggername defaultfontpath ybelkadafonts def checktorchversion if istorchavailable and not istorchgreaterorequalthan111 raise importerror fyou are using torchtorch version but torch1 11 0 is required to use pix2structimageprocessor please upgrade torch adapted from https discuss pytorch orgttfimageextractpatchesinpytorch1714092 def torchextractpatchesimagetensor patchheight patchwidth requiresbackendstorchextractpatches torch checktorchversion imagetensor imagetensor unsqueeze0 patches torch nn functional unfoldimagetensor patchheight patchwidth stridepatchheight patchwidth patches patches reshapeimagetensor size0 imagetensor size1 patchheight patchwidth 1 patches patches permute0 4 2 3 1 reshape imagetensor size2 patchheight imagetensor size3 patchwidth imagetensor size1 patchheight patchwidth return patches unsqueeze0 adapted from https github comgoogleresearchpix2structblob0e1779af0f4db4b652c1d92b3bbd2550a7399123pix2structpreprocessingpreprocessingutils pyl106 def rendertext text str textsize int 36 textcolor str black backgroundcolor str white leftpadding int 5 rightpadding int 5 toppadding int 5 bottompadding int 5 fontbytes optionalbytes none fontpath optionalstr none image image requiresbackendsrendertext vision add new lines so that each line is no more than 80 characters wrapper textwrap textwrapperwidth80 lines wrapper wraptexttext wrappedtext n joinlines if fontbytes is not none and fontpath is none font io bytesiofontbytes elif fontpath is not none font fontpath else font hfhubdownloaddefaultfontpath arial ttf font imagefont truetypefont encodingutf8 sizetextsize use a temporary canvas to determine the width and height in pixels when rendering the text tempdraw imagedraw drawimage newrgb 1 1 backgroundcolor textwidth textheight tempdraw textbbox0 0 wrappedtext font create the actual image with a bit of padding around the text imagewidth textwidth leftpadding rightpadding imageheight textheight toppadding bottompadding image image newrgb imagewidth imageheight backgroundcolor draw imagedraw drawimage draw textxyleftpadding toppadding textwrappedtext filltextcolor fontfont return image adapted from https github comgoogleresearchpix2structblob0e1779af0f4db4b652c1d92b3bbd2550a7399123pix2structpreprocessingpreprocessingutils pyl87 def renderheader image np ndarray header str inputdataformat optionalunionstr childprocesserror none kwargs requiresbackendsrenderheader vision convert to pil image if necessary image topilimageimage inputdataformatinputdataformat headerimage rendertextheader kwargs newwidth maxheaderimage width image width newheight intimage height newwidth image width newheaderheight intheaderimage height newwidth headerimage width newimage image newrgb newwidth newheight newheaderheight white newimage pasteheaderimage resizenewwidth newheaderheight 0 0 newimage pasteimage resizenewwidth newheight 0 newheaderheight convert back to the original framework if necessary newimage tonumpyarraynewimage if inferchanneldimensionformatnewimage channeldimension last newimage tochanneldimensionformatnewimage channeldimension last return newimage class pix2structimageprocessorbaseimageprocessor r constructs a pix2struct image processor args doconvertrgb bool optional defaults to true whether to convert the image to rgb donormalize bool optional defaults to true whether to normalize the image can be overridden by the donormalize parameter in the preprocess method according to pix2struct paper and code the image is normalized with its own mean and standard deviation patchsize dictstr int optional defaults to height 16 width 16 the patch size to use for the image according to pix2struct paper and code the patch size is 16x16 maxpatches int optional defaults to 2048 the maximum number of patches to extract from the image as per the pix2struct paperhttps arxiv orgpdf2210 03347 pdf isvqa bool optional defaults to false whether or not the image processor is for the vqa task if true and headertext is passed in text is rendered onto the input images extract flattened patches from an image args image np ndarray image to extract flattened patches from maxpatches int maximum number of patches to extract patchsize dict dictionary containing the patch height and width returns result np ndarray a sequence of maxpatches flattened patches convert to torch maximize scale s t 1 rows columns patchheight patchwidth imagechannels rows columns patchheight patchwidth imagechannels rows columns 1 offset by 1 so the ids do not contain zeros which represent padding prepare additional patch features rows columns 1 rows columns 2 patchheight patchwidth imagechannels maxpatches 2 patchheight patchwidth imagechannels normalize an image image image imagemean imagestd the image std is to mimic the tensorflow implementation of the perimagestandardization https www tensorflow orgapidocspythontfimageperimagestandardization args image np ndarray image to normalize dataformat str or channeldimension optional the channel dimension format for the output image if unset the channel dimension format of the input image is used inputdataformat str or channeldimension optional the channel dimension format of the input image if not provided it will be inferred take mean across the whole image preprocess an image or batch of images the processor first computes the maximum possible number of aspectratio preserving patches of size patchsize that can be extracted from the image it then pads the image with zeros to make the image respect the constraint of maxpatches before extracting the patches the images are standardized following the tensorflow implementation of perimagestandardization https www tensorflow orgapidocspythontfimageperimagestandardization args images imageinput image to preprocess expects a single or batch of images headertext unionliststr str optional text to render as a header only has an effect if imageprocessor isvqa is true doconvertrgb bool optional defaults to self doconvertrgb whether to convert the image to rgb donormalize bool optional defaults to self donormalize whether to normalize the image maxpatches int optional defaults to self maxpatches maximum number of patches to extract patchsize dict optional defaults to self patchsize dictionary containing the patch height and width 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 channelsfirst or channeldimension first image in numchannels height width format channelslast or channeldimension last image in height width numchannels format unset use 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 pil rgba images are converted to rgb all transformations expect numpy arrays we assume that all images have the same channel dimension format convert to torch tensor and permute create attention mask in numpy 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 pix2struct adapted from https discuss pytorch org t tf image extract patches in pytorch 171409 2 utiliy function to extract patches from a given image tensor returns a tensor of shape 1 patch_height patch_width num_channels x patch_height x patch_width args image_tensor torch tensor the image tensor to extract patches from patch_height int the height of the patches to extract patch_width int the width of the patches to extract adapted from https github com google research pix2struct blob 0e1779af0f4db4b652c1d92b3bbd2550a7399123 pix2struct preprocessing preprocessing_utils py l106 render text this script is entirely adapted from the original script that can be found here https github com google research pix2struct blob main pix2struct preprocessing preprocessing_utils py args text str optional defaults to text to render text_size int optional defaults to 36 size of the text text_color str optional defaults to black color of the text background_color str optional defaults to white color of the background left_padding int optional defaults to 5 padding on the left right_padding int optional defaults to 5 padding on the right top_padding int optional defaults to 5 padding on the top bottom_padding int optional defaults to 5 padding on the bottom font_bytes bytes optional bytes of the font to use if none the default font will be used font_path str optional path to the font to use if none the default font will be used add new lines so that each line is no more than 80 characters use a temporary canvas to determine the width and height in pixels when rendering the text create the actual image with a bit of padding around the text adapted from https github com google research pix2struct blob 0e1779af0f4db4b652c1d92b3bbd2550a7399123 pix2struct preprocessing preprocessing_utils py l87 renders the input text as a header on the input image args image np ndarray the image to render the header on header str the header text data_format union channeldimension str optional the data format of the image can be either channeldimension channels_first or channeldimension channels_last returns np ndarray the image with the header rendered convert to pil image if necessary convert back to the original framework if necessary constructs a pix2struct image processor args do_convert_rgb bool optional defaults to true whether to convert the image to rgb do_normalize bool optional defaults to true whether to normalize the image can be overridden by the do_normalize parameter in the preprocess method according to pix2struct paper and code the image is normalized with its own mean and standard deviation patch_size dict str int optional defaults to height 16 width 16 the patch size to use for the image according to pix2struct paper and code the patch size is 16x16 max_patches int optional defaults to 2048 the maximum number of patches to extract from the image as per the pix2struct paper https arxiv org pdf 2210 03347 pdf is_vqa bool optional defaults to false whether or not the image processor is for the vqa task if true and header_text is passed in text is rendered onto the input images extract flattened patches from an image args image np ndarray image to extract flattened patches from max_patches int maximum number of patches to extract patch_size dict dictionary containing the patch height and width returns result np ndarray a sequence of max_patches flattened patches convert to torch maximize scale s t 1 rows columns patch_height patch_width image_channels rows columns patch_height patch_width image_channels rows columns 1 offset by 1 so the ids do not contain zeros which represent padding prepare additional patch features rows columns 1 rows columns 2 patch_height patch_width image_channels max_patches 2 patch_height patch_width image_channels normalize an image image image image_mean image_std the image std is to mimic the tensorflow implementation of the per_image_standardization https www tensorflow org api_docs python tf image per_image_standardization args image np ndarray image to normalize data_format str or channeldimension optional the channel dimension format for the output image if unset the channel dimension format of the input image is used input_data_format str or channeldimension optional the channel dimension format of the input image if not provided it will be inferred take mean across the whole image preprocess an image or batch of images the processor first computes the maximum possible number of aspect ratio preserving patches of size patch_size that can be extracted from the image it then pads the image with zeros to make the image respect the constraint of max_patches before extracting the patches the images are standardized following the tensorflow implementation of per_image_standardization https www tensorflow org api_docs python tf image per_image_standardization args images imageinput image to preprocess expects a single or batch of images header_text union list str str optional text to render as a header only has an effect if image_processor is_vqa is true do_convert_rgb bool optional defaults to self do_convert_rgb whether to convert the image to rgb do_normalize bool optional defaults to self do_normalize whether to normalize the image max_patches int optional defaults to self max_patches maximum number of patches to extract patch_size dict optional defaults to self patch_size dictionary containing the patch height and width 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 channels_first or channeldimension first image in num_channels height width format channels_last or channeldimension last image in height width num_channels format unset use 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 pil rgba images are converted to rgb all transformations expect numpy arrays we assume that all images have the same channel dimension format convert to torch tensor and permute create attention mask in numpy
import io
import math
from typing import Dict, Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import convert_to_rgb, normalize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    get_image_size,
    infer_channel_dimension_format,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_vision_available, logging
from ...utils.import_utils import requires_backends


if is_vision_available():
    import textwrap

    from PIL import Image, ImageDraw, ImageFont

if is_torch_available():
    import torch

    from transformers.pytorch_utils import is_torch_greater_or_equal_than_1_11
else:
    is_torch_greater_or_equal_than_1_11 = False


logger = logging.get_logger(__name__)
DEFAULT_FONT_PATH = "ybelkada/fonts"


def _check_torch_version():
    if is_torch_available() and not is_torch_greater_or_equal_than_1_11:
        raise ImportError(
            f"You are using torch=={torch.__version__}, but torch>=1.11.0 is required to use "
            "Pix2StructImageProcessor. Please upgrade torch."
        )



def torch_extract_patches(image_tensor, patch_height, patch_width):
    
    requires_backends(torch_extract_patches, ["torch"])
    _check_torch_version()

    image_tensor = image_tensor.unsqueeze(0)
    patches = torch.nn.functional.unfold(image_tensor, (patch_height, patch_width), stride=(patch_height, patch_width))
    patches = patches.reshape(image_tensor.size(0), image_tensor.size(1), patch_height, patch_width, -1)
    patches = patches.permute(0, 4, 2, 3, 1).reshape(
        image_tensor.size(2) // patch_height,
        image_tensor.size(3) // patch_width,
        image_tensor.size(1) * patch_height * patch_width,
    )
    return patches.unsqueeze(0)



def render_text(
    text: str,
    text_size: int = 36,
    text_color: str = "black",
    background_color: str = "white",
    left_padding: int = 5,
    right_padding: int = 5,
    top_padding: int = 5,
    bottom_padding: int = 5,
    font_bytes: Optional[bytes] = None,
    font_path: Optional[str] = None,
) -> Image.Image:
    
    requires_backends(render_text, "vision")
    

    wrapper = textwrap.TextWrapper(width=80)
    lines = wrapper.wrap(text=text)
    wrapped_text = "\n".join(lines)

    if font_bytes is not None and font_path is None:
        font = io.BytesIO(font_bytes)
    elif font_path is not None:
        font = font_path
    else:
        font = hf_hub_download(DEFAULT_FONT_PATH, "Arial.TTF")
    font = ImageFont.truetype(font, encoding="UTF-8", size=text_size)

    
    
    temp_draw = ImageDraw.Draw(Image.new("RGB", (1, 1), background_color))
    _, _, text_width, text_height = temp_draw.textbbox((0, 0), wrapped_text, font)

    
    image_width = text_width + left_padding + right_padding
    image_height = text_height + top_padding + bottom_padding
    image = Image.new("RGB", (image_width, image_height), background_color)
    draw = ImageDraw.Draw(image)
    draw.text(xy=(left_padding, top_padding), text=wrapped_text, fill=text_color, font=font)
    return image



def render_header(
    image: np.ndarray, header: str, input_data_format: Optional[Union[str, ChildProcessError]] = None, **kwargs
):
    
    requires_backends(render_header, "vision")

    
    image = to_pil_image(image, input_data_format=input_data_format)

    header_image = render_text(header, **kwargs)
    new_width = max(header_image.width, image.width)

    new_height = int(image.height * (new_width / image.width))
    new_header_height = int(header_image.height * (new_width / header_image.width))

    new_image = Image.new("RGB", (new_width, new_height + new_header_height), "white")
    new_image.paste(header_image.resize((new_width, new_header_height)), (0, 0))
    new_image.paste(image.resize((new_width, new_height)), (0, new_header_height))

    
    new_image = to_numpy_array(new_image)

    if infer_channel_dimension_format(new_image) == ChannelDimension.LAST:
        new_image = to_channel_dimension_format(new_image, ChannelDimension.LAST)

    return new_image


class Pix2StructImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["flattened_patches"]

    def __init__(
        self,
        do_convert_rgb: bool = True,
        do_normalize: bool = True,
        patch_size: Dict[str, int] = None,
        max_patches: int = 2048,
        is_vqa: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.patch_size = patch_size if patch_size is not None else {"height": 16, "width": 16}
        self.do_normalize = do_normalize
        self.do_convert_rgb = do_convert_rgb
        self.max_patches = max_patches
        self.is_vqa = is_vqa

    def extract_flattened_patches(
        self,
        image: np.ndarray,
        max_patches: int,
        patch_size: dict,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> np.ndarray:
        
        requires_backends(self.extract_flattened_patches, "torch")
        _check_torch_version()

        
        image = to_channel_dimension_format(image, ChannelDimension.FIRST, input_data_format)
        image = torch.from_numpy(image)

        patch_height, patch_width = patch_size["height"], patch_size["width"]
        image_height, image_width = get_image_size(image, ChannelDimension.FIRST)

        
        scale = math.sqrt(max_patches * (patch_height / image_height) * (patch_width / image_width))
        num_feasible_rows = max(min(math.floor(scale * image_height / patch_height), max_patches), 1)
        num_feasible_cols = max(min(math.floor(scale * image_width / patch_width), max_patches), 1)
        resized_height = max(num_feasible_rows * patch_height, 1)
        resized_width = max(num_feasible_cols * patch_width, 1)

        image = torch.nn.functional.interpolate(
            image.unsqueeze(0),
            size=(resized_height, resized_width),
            mode="bilinear",
            align_corners=False,
            antialias=True,
        ).squeeze(0)

        
        patches = torch_extract_patches(image, patch_height, patch_width)

        patches_shape = patches.shape
        rows = patches_shape[1]
        columns = patches_shape[2]
        depth = patches_shape[3]

        
        patches = patches.reshape([rows * columns, depth])

        
        row_ids = torch.arange(rows).reshape([rows, 1]).repeat(1, columns).reshape([rows * columns, 1])
        col_ids = torch.arange(columns).reshape([1, columns]).repeat(rows, 1).reshape([rows * columns, 1])

        
        row_ids += 1
        col_ids += 1

        
        
        row_ids = row_ids.to(torch.float32)
        col_ids = col_ids.to(torch.float32)

        
        result = torch.cat([row_ids, col_ids, patches], -1)

        
        result = torch.nn.functional.pad(result, [0, 0, 0, max_patches - (rows * columns)]).float()

        result = to_numpy_array(result)

        return result

    def normalize(
        self,
        image: np.ndarray,
        data_format: Optional[Union[str, ChannelDimension]] = None,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> np.ndarray:
        
        if image.dtype == np.uint8:
            image = image.astype(np.float32)

        
        mean = np.mean(image)
        std = np.std(image)
        adjusted_stddev = max(std, 1.0 / math.sqrt(np.prod(image.shape)))

        return normalize(
            image,
            mean=mean,
            std=adjusted_stddev,
            data_format=data_format,
            input_data_format=input_data_format,
            **kwargs,
        )

    def preprocess(
        self,
        images: ImageInput,
        header_text: Optional[str] = None,
        do_convert_rgb: bool = None,
        do_normalize: Optional[bool] = None,
        max_patches: Optional[int] = None,
        patch_size: Optional[Dict[str, int]] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: ChannelDimension = ChannelDimension.FIRST,
        input_data_format: Optional[Union[str, ChannelDimension]] = None,
        **kwargs,
    ) -> ImageInput:
        
        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
        patch_size = patch_size if patch_size is not None else self.patch_size
        max_patches = max_patches if max_patches is not None else self.max_patches
        is_vqa = self.is_vqa

        if kwargs.get("data_format", None) is not None:
            raise ValueError("data_format is not an accepted input as the outputs are ")

        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_convert_rgb:
            images = [convert_to_rgb(image) for image in images]

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

        if input_data_format is None:
            
            input_data_format = infer_channel_dimension_format(images[0])

        if is_vqa:
            if header_text is None:
                raise ValueError("A header text must be provided for VQA models.")
            font_bytes = kwargs.pop("font_bytes", None)
            font_path = kwargs.pop("font_path", None)

            if isinstance(header_text, str):
                header_text = [header_text] * len(images)

            images = [
                render_header(image, header_text[i], font_bytes=font_bytes, font_path=font_path)
                for i, image in enumerate(images)
            ]

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

        
        images = [
            self.extract_flattened_patches(
                image=image, max_patches=max_patches, patch_size=patch_size, input_data_format=input_data_format
            )
            for image in images
        ]

        
        attention_masks = [(image.sum(axis=-1) != 0).astype(np.float32) for image in images]

        encoded_outputs = BatchFeature(
            data={"flattened_patches": images, "attention_mask": attention_masks}, tensor_type=return_tensors
        )

        return encoded_outputs