qiMdZddlZddlZddlZddlmZddlmZm Z ddl m Z m Z m Z mZddlmZmZmZmZmZmZmZddlmZdd lmZmZmZmZdd lmZerddl Z dd l!m"Z"m#Z#m$Z$erddl%Z%ejLe'Z(d Z)Gd dedZ*dZ+ d$de,de-de,de,de-de-de-de-de.dzde,dzde"jDfdZ/d%dej`de,d e,e1zdzfd!Z2Gd"d#eZ3d#gZ4y)&z%Image processor class for Pix2Struct.N)hf_hub_download)BaseImageProcessor BatchFeature)convert_to_rgb normalizeto_channel_dimension_format to_pil_image)ChannelDimension ImageInputget_image_sizeinfer_channel_dimension_formatmake_flat_list_of_imagesto_numpy_array valid_images) ImagesKwargs) TensorTypeis_torch_availableis_vision_availablelogging)requires_backends)Image ImageDraw ImageFontzybelkada/fontscVeZdZUdZeed<eeefed<eed<e eezdzed<y)Pix2StructImageProcessorKwargsa max_patches (`int`, *optional*): Maximum number of patches to extract. 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. 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. header_text (`Union[list[str], str]`, *optional*): Text to render as a header. Only has an effect if `image_processor.is_vqa` is `True`. max_patches patch_sizeis_vqaN header_text) __name__ __module__ __qualname____doc__int__annotations__dictstrboollistl/opt/pipecat/venv/lib/python3.12/site-packages/transformers/models/pix2struct/image_processing_pix2struct.pyrr2s5 S#X LcS4''r,rF)totalcttdg|jd}tjj j |||f||f}|j|jd|jd||d}|jdddddj|jd|z|jd|z|jd|z|z}|jdS) a Utility function to extract patches from a given image tensor. Returns a tensor of shape (1, `rows`, `columns`, `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. torchr)strider) rtorch_extract_patches unsqueezer0nn functionalunfoldreshapesizepermute) image_tensor patch_height patch_widthpatchess r-r6r6Fs+gY7))!,Lhh!!(( k7R\hju[v(wGool//2L4E4Ea4H,XceghGooaAq!,44! ,! +!|+k9G   Q r,text text_size text_colorbackground_color left_padding right_padding top_paddingbottom_padding font_bytes font_pathreturnc Tttdtjd} | j |} dj | } || t j|} n| | } nttd} tj| d|} tjtjd d |}|j!d | | \}}}}||z|z}||z|z}tjd ||f|}tj|}|j#||f| || |S) a 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. visionP)width)rB z Arial.TTFzUTF-8)encodingr<RGB)r2r2rr)xyrBfillfont)r render_texttextwrap TextWrapperwrapjoinioBytesIOrDEFAULT_FONT_PATHrtruetyperDrawrnewtextbboxrB)rBrCrDrErFrGrHrIrJrKwrapperlines wrapped_textrW temp_draw_ text_width text_height image_width image_heightimagedraws r-rXrXas$Lk8,"",G LLdL #E99U#L)"3zz*%  0+>   dW9 EDuyy8HIJI$-$6$6v|T$R!Aq*k|+m;K,~=L IIek<8:J KE >>% DII, ,JL//9|?Q?Q3QRS %)Z:K-K!LgVI OOL''4E(FGP OOELL)Z!89A?P;QRy)I%i04D4I4II/ ;K;P;PQ r,ceZdZdZddgZeZ ddededee e fdzde d ed df fd Z dd e jde ded e ezdzd e jf dZ dd e jde ezdzd e ezdzd e jfdZddddddej"dfdede dzdedzdedzde dzdee e fdzde ezdzded e ezdzd efdZxZS)Pix2StructImageProcessoraj 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://huggingface.co/papers/2210.03347). 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. flattened_patchesattention_maskNdo_convert_rgb do_normalizerrrrLc xt|di|||nddd|_||_||_||_||_y)N)rurPr+)super__init__rrrrr)selfrrrrrry __class__s r-rz!Pix2StructImageProcessor.__init__sH "6"(2(>*r\^D_(,& r,rmrpc t|jdt|tj|}t j |}|d|d}}t|tj\}} tj|||z z|| z z} tttj| |z|z |d} tttj| | z|z |d} t| |zd} t| |zd}t jjj|j!d| |fddd j#d}t%|||}|j&}|d}|d }|d }|j)||z|g}t j*|j)|dgj-d|j)||zdg}t j*|j)d|gj-|dj)||zdg}|dz }|dz }|j/t j0}|j/t j0}t j2|||gd }t jjj5|ddd|||zz gj7}t9|}|S) a 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. r0rurPr2rbilinearFT)r<mode align_corners antialiasr5rr3)rextract_flattened_patchesr r FIRSTr0 from_numpyr mathsqrtrtminfloorr8r9 interpolater7squeezer6shaper;arangerepeattofloat32catpadfloatr)rrmrrrpryr?r@rlrkscalenum_feasible_rowsnum_feasible_colsresized_height resized_widthrA patches_shaperowscolumnsdepthrow_idscol_idsresults r-rz2Pix2StructImageProcessor.extract_flattened_patchess. $88'B,E3C3I3IK\]  '$.x$8*W:Mk $25:J:P:P$Q! k + )DEWbIbcdDJJu|/Cl/R$SU` acdeDJJu{/B[/P$QS^ _abc.=qA- ;Q? ##// OOA  -0 0 '!* (|[I Q"a //4'>5"9:,,t$,,dAY7>>q'JRRTX[bTbdeSfg,,w'//G =DDT1MUUW[^eWeghVij 1 1 **U]]+**U]]+GWg6;$$((!Q;$QX.;Y1Z[aac' r, data_formatc n|jtjk(r|jtj}tj |}tj |}t|dtjtj|jz }t|f||||d|S)a- Normalize an image. image = (image - image_mean) / image_std. 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. g?)meanstdrrp) dtypenpuint8astyperrrrtrrprodrr)rrmrrpryrradjusted_stddevs r-rz"Pix2StructImageProcessor.normalizeFs& ;;"(( "LL,Ewwu~ffUmc32775;;3G)H#HI  #/     r,imagesr return_tensorsc ||n |j}||n |j}||n |j}||n |j}|j} | j d t dt|}t|s t d|r|D cgc] } t| }} |D cgc] } t| }} | t|d} | r}| t d| jdd} | jdd}t|tr|gt|z}t!|D cgc]\}} t#| ||| | }}} |r |D cgc]} |j%| | }} |D cgc]} |j'| ||| }} |D cgc]4} | j)d dk7j+t,j.6}} t1||d|}|Scc} wcc} wcc} }wcc} wcc} wcc} w)a 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`. 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.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`. - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.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. Nrz8data_format is not an accepted input as the outputs are zSInvalid image type. Must be of type PIL.Image.Image, numpy.ndarray, or torch.Tensorrz.A header text must be provided for VQA models.rJrK)rJrK)rmrp)rmrrrpr3)axis)rr)data tensor_type)rrrrrget ValueErrorrrrrrpop isinstancer(len enumeratersrrsumrrrr)rrr rrrrrrrpryrrmrJrKiattention_masksencoded_outputss r- preprocessz#Pix2StructImageProcessor.preprocessjs7`(4'?|TEVEV +9+E4K^K^#-#9Zt %0%vay I  " !QRRL$7J ;5I+s+*mc&k9 !*& 1Aue[^ V_`F djk[`dnn5DUnVkFk    * *_p +   V\\EEII2I.!3;;BJJG\\&'-Q_m SA= l ]s$G (G%,G*G01G59G:)TTNiFN)NN)r!r"r#r$model_input_namesr valid_kwargsr)r'r(r%rrndarrayr rrrr rr __classcell__)rs@r-rrs(-.>?1L $!,0 cNT)      *<@ OzzOO O !11D8 O Oh6:;? " zz" ++d2" !11D8 "  " N#'&*$("&,026(8(>(>;?ii4Zit i Tk i 4Z icNT)ij(4/i&i!11D8i ir,r) $blackrrrrrNNr)5r$r]rnumpyrhuggingface_hubrimage_processing_utilsrrimage_transformsrrr r image_utilsr r r rrrrprocessing_utilsrutilsrrrrutils.import_utilsrrYPILrrrr0 get_loggerr!loggerr_rr6r(r%bytesrXrChildProcessErrorrsr__all__r+r,r-rsk, +Fdd-QQ3//   H %$(\(( :## @ @@@ @  @  @@@ @Tz@ [[@H%%S%SK\E\_cEc%PF1FR & &r,