Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2022 The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Image processor class for LayoutLMv2."""16 17from typing import Optional, Union18 19import numpy as np20 21from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict22from ...image_transforms import flip_channel_order, resize, to_channel_dimension_format, to_pil_image23from ...image_utils import (24    ChannelDimension,25    ImageInput,26    PILImageResampling,27    infer_channel_dimension_format,28    make_flat_list_of_images,29    to_numpy_array,30    valid_images,31    validate_preprocess_arguments,32)33from ...utils import (34    TensorType,35    filter_out_non_signature_kwargs,36    is_pytesseract_available,37    is_vision_available,38    logging,39    requires_backends,40)41from ...utils.import_utils import requires42 43 44if is_vision_available():45    import PIL46 47# soft dependency48if is_pytesseract_available():49    import pytesseract50 51logger = logging.get_logger(__name__)52 53 54def normalize_box(box, width, height):55    return [56        int(1000 * (box[0] / width)),57        int(1000 * (box[1] / height)),58        int(1000 * (box[2] / width)),59        int(1000 * (box[3] / height)),60    ]61 62 63def apply_tesseract(64    image: np.ndarray,65    lang: Optional[str],66    tesseract_config: Optional[str] = None,67    input_data_format: Optional[Union[str, ChannelDimension]] = None,68):69    """Applies Tesseract OCR on a document image, and returns recognized words + normalized bounding boxes."""70    tesseract_config = tesseract_config if tesseract_config is not None else ""71 72    # apply OCR73    pil_image = to_pil_image(image, input_data_format=input_data_format)74    image_width, image_height = pil_image.size75    data = pytesseract.image_to_data(pil_image, lang=lang, output_type="dict", config=tesseract_config)76    words, left, top, width, height = data["text"], data["left"], data["top"], data["width"], data["height"]77 78    # filter empty words and corresponding coordinates79    irrelevant_indices = [idx for idx, word in enumerate(words) if not word.strip()]80    words = [word for idx, word in enumerate(words) if idx not in irrelevant_indices]81    left = [coord for idx, coord in enumerate(left) if idx not in irrelevant_indices]82    top = [coord for idx, coord in enumerate(top) if idx not in irrelevant_indices]83    width = [coord for idx, coord in enumerate(width) if idx not in irrelevant_indices]84    height = [coord for idx, coord in enumerate(height) if idx not in irrelevant_indices]85 86    # turn coordinates into (left, top, left+width, top+height) format87    actual_boxes = []88    for x, y, w, h in zip(left, top, width, height):89        actual_box = [x, y, x + w, y + h]90        actual_boxes.append(actual_box)91 92    # finally, normalize the bounding boxes93    normalized_boxes = []94    for box in actual_boxes:95        normalized_boxes.append(normalize_box(box, image_width, image_height))96 97    assert len(words) == len(normalized_boxes), "Not as many words as there are bounding boxes"98 99    return words, normalized_boxes100 101 102@requires(backends=("vision",))103class LayoutLMv2ImageProcessor(BaseImageProcessor):104    r"""105    Constructs a LayoutLMv2 image processor.106 107    Args:108        do_resize (`bool`, *optional*, defaults to `True`):109            Whether to resize the image's (height, width) dimensions to `(size["height"], size["width"])`. Can be110            overridden by `do_resize` in `preprocess`.111        size (`dict[str, int]` *optional*, defaults to `{"height": 224, "width": 224}`):112            Size of the image after resizing. Can be overridden by `size` in `preprocess`.113        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BILINEAR`):114            Resampling filter to use if resizing the image. Can be overridden by the `resample` parameter in the115            `preprocess` method.116        apply_ocr (`bool`, *optional*, defaults to `True`):117            Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by118            `apply_ocr` in `preprocess`.119        ocr_lang (`str`, *optional*):120            The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is121            used. Can be overridden by `ocr_lang` in `preprocess`.122        tesseract_config (`str`, *optional*, defaults to `""`):123            Any additional custom configuration flags that are forwarded to the `config` parameter when calling124            Tesseract. For example: '--psm 6'. Can be overridden by `tesseract_config` in `preprocess`.125    """126 127    model_input_names = ["pixel_values"]128 129    def __init__(130        self,131        do_resize: bool = True,132        size: Optional[dict[str, int]] = None,133        resample: PILImageResampling = PILImageResampling.BILINEAR,134        apply_ocr: bool = True,135        ocr_lang: Optional[str] = None,136        tesseract_config: Optional[str] = "",137        **kwargs,138    ) -> None:139        super().__init__(**kwargs)140        size = size if size is not None else {"height": 224, "width": 224}141        size = get_size_dict(size)142 143        self.do_resize = do_resize144        self.size = size145        self.resample = resample146        self.apply_ocr = apply_ocr147        self.ocr_lang = ocr_lang148        self.tesseract_config = tesseract_config149 150    # Copied from transformers.models.vit.image_processing_vit.ViTImageProcessor.resize151    def resize(152        self,153        image: np.ndarray,154        size: dict[str, int],155        resample: PILImageResampling = PILImageResampling.BILINEAR,156        data_format: Optional[Union[str, ChannelDimension]] = None,157        input_data_format: Optional[Union[str, ChannelDimension]] = None,158        **kwargs,159    ) -> np.ndarray:160        """161        Resize an image to `(size["height"], size["width"])`.162 163        Args:164            image (`np.ndarray`):165                Image to resize.166            size (`dict[str, int]`):167                Dictionary in the format `{"height": int, "width": int}` specifying the size of the output image.168            resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`):169                `PILImageResampling` filter to use when resizing the image e.g. `PILImageResampling.BILINEAR`.170            data_format (`ChannelDimension` or `str`, *optional*):171                The channel dimension format for the output image. If unset, the channel dimension format of the input172                image is used. Can be one of:173                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.174                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.175                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.176            input_data_format (`ChannelDimension` or `str`, *optional*):177                The channel dimension format for the input image. If unset, the channel dimension format is inferred178                from the input image. Can be one of:179                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.180                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.181                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.182 183        Returns:184            `np.ndarray`: The resized image.185        """186        size = get_size_dict(size)187        if "height" not in size or "width" not in size:188            raise ValueError(f"The `size` dictionary must contain the keys `height` and `width`. Got {size.keys()}")189        output_size = (size["height"], size["width"])190        return resize(191            image,192            size=output_size,193            resample=resample,194            data_format=data_format,195            input_data_format=input_data_format,196            **kwargs,197        )198 199    @filter_out_non_signature_kwargs()200    def preprocess(201        self,202        images: ImageInput,203        do_resize: Optional[bool] = None,204        size: Optional[dict[str, int]] = None,205        resample: Optional[PILImageResampling] = None,206        apply_ocr: Optional[bool] = None,207        ocr_lang: Optional[str] = None,208        tesseract_config: Optional[str] = None,209        return_tensors: Optional[Union[str, TensorType]] = None,210        data_format: ChannelDimension = ChannelDimension.FIRST,211        input_data_format: Optional[Union[str, ChannelDimension]] = None,212    ) -> PIL.Image.Image:213        """214        Preprocess an image or batch of images.215 216        Args:217            images (`ImageInput`):218                Image to preprocess.219            do_resize (`bool`, *optional*, defaults to `self.do_resize`):220                Whether to resize the image.221            size (`dict[str, int]`, *optional*, defaults to `self.size`):222                Desired size of the output image after resizing.223            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):224                Resampling filter to use if resizing the image. This can be one of the enum `PIL.Image` resampling225                filter. Only has an effect if `do_resize` is set to `True`.226            apply_ocr (`bool`, *optional*, defaults to `self.apply_ocr`):227                Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes.228            ocr_lang (`str`, *optional*, defaults to `self.ocr_lang`):229                The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is230                used.231            tesseract_config (`str`, *optional*, defaults to `self.tesseract_config`):232                Any additional custom configuration flags that are forwarded to the `config` parameter when calling233                Tesseract.234            return_tensors (`str` or `TensorType`, *optional*):235                The type of tensors to return. Can be one of:236                    - Unset: Return a list of `np.ndarray`.237                    - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.238                    - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.239                    - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.240                    - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.241            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):242                The channel dimension format for the output image. Can be one of:243                    - `ChannelDimension.FIRST`: image in (num_channels, height, width) format.244                    - `ChannelDimension.LAST`: image in (height, width, num_channels) format.245        """246        do_resize = do_resize if do_resize is not None else self.do_resize247        size = size if size is not None else self.size248        size = get_size_dict(size)249        resample = resample if resample is not None else self.resample250        apply_ocr = apply_ocr if apply_ocr is not None else self.apply_ocr251        ocr_lang = ocr_lang if ocr_lang is not None else self.ocr_lang252        tesseract_config = tesseract_config if tesseract_config is not None else self.tesseract_config253 254        images = make_flat_list_of_images(images)255 256        if not valid_images(images):257            raise ValueError(258                "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "259                "torch.Tensor, tf.Tensor or jax.ndarray."260            )261        validate_preprocess_arguments(262            do_resize=do_resize,263            size=size,264            resample=resample,265        )266 267        # All transformations expect numpy arrays.268        images = [to_numpy_array(image) for image in images]269 270        if input_data_format is None:271            # We assume that all images have the same channel dimension format.272            input_data_format = infer_channel_dimension_format(images[0])273 274        if apply_ocr:275            requires_backends(self, "pytesseract")276            words_batch = []277            boxes_batch = []278            for image in images:279                words, boxes = apply_tesseract(image, ocr_lang, tesseract_config, input_data_format=input_data_format)280                words_batch.append(words)281                boxes_batch.append(boxes)282 283        if do_resize:284            images = [285                self.resize(image=image, size=size, resample=resample, input_data_format=input_data_format)286                for image in images287            ]288 289        # flip color channels from RGB to BGR (as Detectron2 requires this)290        images = [flip_channel_order(image, input_data_format=input_data_format) for image in images]291        images = [292            to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format) for image in images293        ]294 295        data = BatchFeature(data={"pixel_values": images}, tensor_type=return_tensors)296 297        if apply_ocr:298            data["words"] = words_batch299            data["boxes"] = boxes_batch300        return data301 302 303__all__ = ["LayoutLMv2ImageProcessor"]304