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 layoutlmv2 from typing import dict optional union import numpy as np from imageprocessingutils import baseimageprocessor batchfeature getsizedict from imagetransforms import flipchannelorder resize tochanneldimensionformat topilimage from imageutils import channeldimension imageinput pilimageresampling inferchanneldimensionformat makelistofimages tonumpyarray validimages from utils import tensortype ispytesseractavailable isvisionavailable logging requiresbackends if isvisionavailable import pil soft dependency if ispytesseractavailable import pytesseract logger logging getloggername def normalizeboxbox width height return int1000 box0 width int1000 box1 height int1000 box2 width int1000 box3 height def applytesseract image np ndarray lang optionalstr tesseractconfig optionalstr none inputdataformat optionalunionstr channeldimension none apply ocr filter empty words and corresponding coordinates turn coordinates into left top leftwidth topheight format finally normalize the bounding boxes modelinputnames pixelvalues def init self doresize bool true size dictstr int none resample pilimageresampling pilimageresampling bilinear applyocr bool true ocrlang optionalstr none tesseractconfig optionalstr kwargs none super initkwargs size size if size is not none else height 224 width 224 size getsizedictsize self doresize doresize self size size self resample resample self applyocr applyocr self ocrlang ocrlang self tesseractconfig tesseractconfig copied from transformers models vit imageprocessingvit vitimageprocessor resize def resize self image np ndarray size dictstr int resample pilimageresampling pilimageresampling bilinear dataformat optionalunionstr channeldimension none inputdataformat optionalunionstr channeldimension none kwargs np ndarray size getsizedictsize if height not in size or width not in size raise valueerrorfthe size dictionary must contain the keys height and width got size keys outputsize sizeheight sizewidth return resize image sizeoutputsize resampleresample dataformatdataformat inputdataformatinputdataformat kwargs def preprocess self images imageinput doresize bool none size dictstr int none resample pilimageresampling none applyocr bool none ocrlang optionalstr none tesseractconfig optionalstr none returntensors optionalunionstr tensortype none dataformat channeldimension channeldimension first inputdataformat optionalunionstr channeldimension none kwargs pil image image doresize doresize if doresize is not none else self doresize size size if size is not none else self size size getsizedictsize resample resample if resample is not none else self resample applyocr applyocr if applyocr is not none else self applyocr ocrlang ocrlang if ocrlang is not none else self ocrlang tesseractconfig tesseractconfig if tesseractconfig is not none else self tesseractconfig images makelistofimagesimages if not validimagesimages raise valueerror invalid image type must be of type pil image image numpy ndarray torch tensor tf tensor or jax ndarray if doresize and size is none raise valueerrorsize must be specified if doresize is true all transformations expect numpy arrays images tonumpyarrayimage for image in images if inputdataformat is none we assume that all images have the same channel dimension format inputdataformat inferchanneldimensionformatimages0 if applyocr requiresbackendsself pytesseract wordsbatch boxesbatch for image in images words boxes applytesseractimage ocrlang tesseractconfig inputdataformatinputdataformat wordsbatch appendwords boxesbatch appendboxes if doresize images self resizeimageimage sizesize resampleresample inputdataformatinputdataformat for image in images flip color channels from rgb to bgr as detectron2 requires this images flipchannelorderimage inputdataformatinputdataformat for image in images images tochanneldimensionformatimage dataformat inputchanneldiminputdataformat for image in images data batchfeaturedatapixelvalues images tensortypereturntensors if applyocr datawords wordsbatch databoxes boxesbatch return data 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 layoutlmv2 soft dependency applies tesseract ocr on a document image and returns recognized words normalized bounding boxes apply ocr filter empty words and corresponding coordinates turn coordinates into left top left width top height format finally normalize the bounding boxes constructs a layoutlmv2 image processor args do_resize bool optional defaults to true whether to resize the image s height width dimensions to size height size width can be overridden by do_resize in preprocess size dict str int optional defaults to height 224 width 224 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 the resample parameter in the preprocess method apply_ocr bool optional defaults to true whether to apply the tesseract ocr engine to get words normalized bounding boxes can be overridden by apply_ocr in preprocess ocr_lang str optional the language specified by its iso code to be used by the tesseract ocr engine by default english is used can be overridden by ocr_lang in preprocess tesseract_config str optional defaults to any additional custom configuration flags that are forwarded to the config parameter when calling tesseract for example psm 6 can be overridden by tesseract_config in preprocess 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 preprocess an image or batch of images args images imageinput image to preprocess do_resize bool optional defaults to self do_resize whether to resize the image size dict str int optional defaults to self size desired size of the output image after resizing resample pilimageresampling optional defaults to self resample resampling filter to use if resizing the image this can be one of the enum pil image resampling filter only has an effect if do_resize is set to true apply_ocr bool optional defaults to self apply_ocr whether to apply the tesseract ocr engine to get words normalized bounding boxes ocr_lang str optional defaults to self ocr_lang the language specified by its iso code to be used by the tesseract ocr engine by default english is used tesseract_config str optional defaults to self tesseract_config any additional custom configuration flags that are forwarded to the config parameter when calling tesseract 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 all transformations expect numpy arrays we assume that all images have the same channel dimension format flip color channels from rgb to bgr as detectron2 requires this
from typing import Dict, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import flip_channel_order, resize, to_channel_dimension_format, to_pil_image
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    infer_channel_dimension_format,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_pytesseract_available, is_vision_available, logging, requires_backends


if is_vision_available():
    import PIL


if is_pytesseract_available():
    import pytesseract

logger = logging.get_logger(__name__)


def normalize_box(box, width, height):
    return [
        int(1000 * (box[0] / width)),
        int(1000 * (box[1] / height)),
        int(1000 * (box[2] / width)),
        int(1000 * (box[3] / height)),
    ]


def apply_tesseract(
    image: np.ndarray,
    lang: Optional[str],
    tesseract_config: Optional[str] = None,
    input_data_format: Optional[Union[str, ChannelDimension]] = None,
):
    
    tesseract_config = tesseract_config if tesseract_config is not None else ""

    
    pil_image = to_pil_image(image, input_data_format=input_data_format)
    image_width, image_height = pil_image.size
    data = pytesseract.image_to_data(pil_image, lang=lang, output_type="dict", config=tesseract_config)
    words, left, top, width, height = data["text"], data["left"], data["top"], data["width"], data["height"]

    
    irrelevant_indices = [idx for idx, word in enumerate(words) if not word.strip()]
    words = [word for idx, word in enumerate(words) if idx not in irrelevant_indices]
    left = [coord for idx, coord in enumerate(left) if idx not in irrelevant_indices]
    top = [coord for idx, coord in enumerate(top) if idx not in irrelevant_indices]
    width = [coord for idx, coord in enumerate(width) if idx not in irrelevant_indices]
    height = [coord for idx, coord in enumerate(height) if idx not in irrelevant_indices]

    
    actual_boxes = []
    for x, y, w, h in zip(left, top, width, height):
        actual_box = [x, y, x + w, y + h]
        actual_boxes.append(actual_box)

    
    normalized_boxes = []
    for box in actual_boxes:
        normalized_boxes.append(normalize_box(box, image_width, image_height))

    assert len(words) == len(normalized_boxes), "Not as many words as there are bounding boxes"

    return words, normalized_boxes


class LayoutLMv2ImageProcessor(BaseImageProcessor):
    r

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        do_resize: bool = True,
        size: Dict[str, int] = None,
        resample: PILImageResampling = PILImageResampling.BILINEAR,
        apply_ocr: bool = True,
        ocr_lang: Optional[str] = None,
        tesseract_config: Optional[str] = "",
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        size = size if size is not None else {"height": 224, "width": 224}
        size = get_size_dict(size)

        self.do_resize = do_resize
        self.size = size
        self.resample = resample
        self.apply_ocr = apply_ocr
        self.ocr_lang = ocr_lang
        self.tesseract_config = tesseract_config

    
    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 preprocess(
        self,
        images: ImageInput,
        do_resize: bool = None,
        size: Dict[str, int] = None,
        resample: PILImageResampling = None,
        apply_ocr: bool = None,
        ocr_lang: Optional[str] = None,
        tesseract_config: Optional[str] = None,
        return_tensors: Optional[Union[str, TensorType]] = None,
        data_format: 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
        apply_ocr = apply_ocr if apply_ocr is not None else self.apply_ocr
        ocr_lang = ocr_lang if ocr_lang is not None else self.ocr_lang
        tesseract_config = tesseract_config if tesseract_config is not None else self.tesseract_config

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

        
        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 apply_ocr:
            requires_backends(self, "pytesseract")
            words_batch = []
            boxes_batch = []
            for image in images:
                words, boxes = apply_tesseract(image, ocr_lang, tesseract_config, input_data_format=input_data_format)
                words_batch.append(words)
                boxes_batch.append(boxes)

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

        
        images = [flip_channel_order(image, 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 = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)

        if apply_ocr:
            data["words"] = words_batch
            data["boxes"] = boxes_batch
        return data