Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 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"""Fast Image processor class for LayoutLMv2."""16 17from typing import Optional, Union18 19import torch20from torchvision.transforms.v2 import functional as F21 22from ...image_processing_utils_fast import BaseImageProcessorFast, BatchFeature, DefaultFastImageProcessorKwargs23from ...image_transforms import ChannelDimension, group_images_by_shape, reorder_images24from ...image_utils import ImageInput, PILImageResampling, SizeDict25from ...processing_utils import Unpack26from ...utils import (27    TensorType,28    auto_docstring,29    logging,30    requires_backends,31)32from .image_processing_layoutlmv2 import apply_tesseract33 34 35logger = logging.get_logger(__name__)36 37 38class LayoutLMv2FastImageProcessorKwargs(DefaultFastImageProcessorKwargs):39    """40    Args:41        apply_ocr (`bool`, *optional*, defaults to `True`):42            Whether to apply the Tesseract OCR engine to get words + normalized bounding boxes. Can be overridden by43            the `apply_ocr` parameter in the `preprocess` method.44        ocr_lang (`str`, *optional*):45            The language, specified by its ISO code, to be used by the Tesseract OCR engine. By default, English is46            used. Can be overridden by the `ocr_lang` parameter in the `preprocess` method.47        tesseract_config (`str`, *optional*):48            Any additional custom configuration flags that are forwarded to the `config` parameter when calling49            Tesseract. For example: '--psm 6'. Can be overridden by the `tesseract_config` parameter in the50            `preprocess` method.51    """52 53    apply_ocr: Optional[bool]54    ocr_lang: Optional[str]55    tesseract_config: Optional[str]56 57 58@auto_docstring59class LayoutLMv2ImageProcessorFast(BaseImageProcessorFast):60    resample = PILImageResampling.BILINEAR61    size = {"height": 224, "width": 224}62    rescale_factor = None63    do_resize = True64    apply_ocr = True65    ocr_lang = None66    tesseract_config = ""67    valid_kwargs = LayoutLMv2FastImageProcessorKwargs68 69    def __init__(self, **kwargs: Unpack[LayoutLMv2FastImageProcessorKwargs]):70        super().__init__(**kwargs)71 72    @auto_docstring73    def preprocess(self, images: ImageInput, **kwargs: Unpack[LayoutLMv2FastImageProcessorKwargs]) -> BatchFeature:74        return super().preprocess(images, **kwargs)75 76    def _preprocess(77        self,78        images: list["torch.Tensor"],79        do_resize: bool,80        size: SizeDict,81        interpolation: Optional["F.InterpolationMode"],82        apply_ocr: bool,83        ocr_lang: Optional[str],84        tesseract_config: Optional[str],85        disable_grouping: Optional[bool],86        return_tensors: Optional[Union[str, TensorType]],87        **kwargs,88    ) -> BatchFeature:89        # Tesseract OCR to get words + normalized bounding boxes90        if apply_ocr:91            requires_backends(self, "pytesseract")92            words_batch = []93            boxes_batch = []94            for image in images:95                if image.is_cuda:96                    logger.warning_once(97                        "apply_ocr can only be performed on cpu. Tensors will be transferred to cpu before processing."98                    )99                words, boxes = apply_tesseract(100                    image.cpu(), ocr_lang, tesseract_config, input_data_format=ChannelDimension.FIRST101                )102                words_batch.append(words)103                boxes_batch.append(boxes)104 105        # Group images by size for batched resizing106        grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)107        resized_images_grouped = {}108        for shape, stacked_images in grouped_images.items():109            if do_resize:110                stacked_images = self.resize(image=stacked_images, size=size, interpolation=interpolation)111            resized_images_grouped[shape] = stacked_images112        resized_images = reorder_images(resized_images_grouped, grouped_images_index)113 114        # Group images by size for further processing115        # Needed in case do_resize is False, or resize returns images with different sizes116        grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)117        processed_images_grouped = {}118        for shape, stacked_images in grouped_images.items():119            # flip color channels from RGB to BGR (as Detectron2 requires this)120            stacked_images = stacked_images.flip(1)121            processed_images_grouped[shape] = stacked_images122 123        processed_images = reorder_images(processed_images_grouped, grouped_images_index)124        processed_images = torch.stack(processed_images, dim=0) if return_tensors else processed_images125 126        data = BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)127 128        if apply_ocr:129            data["words"] = words_batch130            data["boxes"] = boxes_batch131 132        return data133 134 135__all__ = ["LayoutLMv2ImageProcessorFast"]136