Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2024 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 16 17from typing import Optional, Union18 19import numpy as np20 21from transformers.processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack22from transformers.tokenization_utils_base import PreTokenizedInput, TextInput23 24from ...image_processing_utils import BatchFeature25from ...image_utils import ImageInput26from ...utils import is_vision_available, logging27 28 29if is_vision_available():30    from ...image_utils import load_images31 32logger = logging.get_logger(__name__)33 34 35class GotOcr2TextKwargs(TextKwargs, total=False):36    format: Optional[bool]37 38 39class GotOcr2ImagesKwargs(ImagesKwargs, total=False):40    box: Optional[Union[list, tuple[float, float], tuple[float, float, float, float]]]41    color: Optional[str]42    num_image_tokens: Optional[int]43    multi_page: Optional[bool]44    crop_to_patches: Optional[bool]45    min_patches: Optional[int]46    max_patches: Optional[int]47 48 49class GotOcr2ProcessorKwargs(ProcessingKwargs, total=False):50    text_kwargs: GotOcr2TextKwargs51    images_kwargs: GotOcr2ImagesKwargs52    _defaults = {53        "text_kwargs": {54            "padding": False,55            "format": False,56        },57        "images_kwargs": {58            "num_image_tokens": 256,59            "multi_page": False,60            "crop_to_patches": False,61            "min_patches": 1,62            "max_patches": 12,63        },64    }65 66 67def preprocess_box_annotation(box: Union[list, tuple], image_size: tuple[int, int]) -> list:68    """69    Convert box annotation to the format [x1, y1, x2, y2] in the range [0, 1000].70    """71    width, height = image_size72    if len(box) == 4:73        box[0] = int(box[0] / width * 1000)74        box[1] = int(box[1] / height * 1000)75        box[2] = int(box[2] / width * 1000)76        box[3] = int(box[3] / height * 1000)77    else:78        raise ValueError("Box must be a list or tuple of lists in the form [x1, y1, x2, y2].")79 80    return list(box)81 82 83class GotOcr2Processor(ProcessorMixin):84    r"""85    Constructs a GotOcr2 processor which wraps a [`GotOcr2ImageProcessor`] and86    [`PretrainedTokenizerFast`] tokenizer into a single processor that inherits both the image processor and87    tokenizer functionalities. See the [`~GotOcr2Processor.__call__`] and [`~GotOcr2Processor.decode`] for more information.88    Args:89        image_processor ([`GotOcr2ImageProcessor`], *optional*):90            The image processor is a required input.91        tokenizer ([`PreTrainedTokenizer`, `PreTrainedTokenizerFast`], *optional*):92            The tokenizer is a required input.93        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages94            in a chat into a tokenizable string.95    """96 97    attributes = ["image_processor", "tokenizer"]98    image_processor_class = "AutoImageProcessor"99    tokenizer_class = "PreTrainedTokenizerFast"100 101    def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):102        super().__init__(image_processor, tokenizer, chat_template=chat_template)103 104        self.message_start_token = "<|im_start|>"105        self.message_end_token = "<|im_end|>"106        self.img_start_token = "<img>"107        self.img_end_token = "</img>"108        self.img_pad_token = "<imgpad>"109        self.image_token = "<imgpad>"  # keep the above for BC, but we need to call it `image_token`110        self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)111        self.system_query = "system\nYou should follow the instructions carefully and explain your answers in detail."112 113    def _make_list_of_inputs(self, images, text, box, color, multi_page):114        if not isinstance(images, (list, tuple)):115            images = [images]116            if multi_page:117                logger.warning("Multi-page inference is enabled but only one image is passed.")118                images = [images]119        elif isinstance(images[0], (list, tuple)) and not multi_page:120            raise ValueError("Nested images are only supported with `multi_page` set to `True`.")121        elif not isinstance(images[0], (list, tuple)) and multi_page:122            images = [images]123 124        if isinstance(text, str):125            text = [text]126 127        if not isinstance(box[0], (list, tuple)):128            # Use the same box for all images129            box = [box for _ in range(len(images))]130        if not isinstance(color, (list, tuple)):131            color = [color for _ in range(len(images))]132 133        return images, text, box, color134 135    def __call__(136        self,137        images: Optional[ImageInput] = None,138        text: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] = None,139        audio=None,140        videos=None,141        **kwargs: Unpack[GotOcr2ProcessorKwargs],142    ) -> BatchFeature:143        """144        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`145        and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] to encode the text if `text`146        is not `None`, otherwise encode default OCR queries which depends on the `format`, `box`, `color`, `multi_page` and147        `crop_to_patches` arguments. To prepare the vision inputs, this method forwards the `images` and `kwargs` arguments to148        GotOcr2ImageProcessor's [`~GotOcr2ImageProcessor.__call__`] if `images` is not `None`.149 150        Args:151            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):152                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch153                tensor. Both channels-first and channels-last formats are supported.154            text (`str`, `list[str]`, `list[list[str]]`):155                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings156                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set157                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).158            format (`bool`, *optional*):159                If set, will add the format token to the query, and the model will return the OCR result with formatting.160            box (`list[float]`, `list[tuple[float, float]]`, `list[tuple[float, float, float, float]]`, *optional*):161                The box annotation to be added to the query. If a list of floats or a tuple of floats is provided, it162                will be interpreted as [x1, y1, x2, y2]. If a list of tuples is provided, each tuple should be in the163                form (x1, y1, x2, y2).164            color (`str`, *optional*):165                The color annotation to be added to the query. The model will return the OCR result within the box with166                the specified color.167            multi_page (`bool`, *optional*):168                If set, will enable multi-page inference. The model will return the OCR result across multiple pages.169            crop_to_patches (`bool`, *optional*):170                If set, will crop the image to patches. The model will return the OCR result upon the patch reference.171            min_patches (`int`, *optional*):172                The minimum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to173                `True`.174            max_patches (`int`, *optional*):175                The maximum number of patches to be cropped from the image. Only used when `crop_to_patches` is set to176                `True`.177 178            return_tensors (`str` or [`~utils.TensorType`], *optional*):179                If set, will return tensors of a particular framework. Acceptable values are:180                - `'tf'`: Return TensorFlow `tf.constant` objects.181                - `'pt'`: Return PyTorch `torch.Tensor` objects.182                - `'np'`: Return NumPy `np.ndarray` objects.183                - `'jax'`: Return JAX `jnp.ndarray` objects.184 185        Returns:186            [`BatchFeature`]: A [`BatchFeature`] with the following fields:187 188            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.189            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when190              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not191              `None`).192            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.193        """194 195        output_kwargs = self._merge_kwargs(196            GotOcr2ProcessorKwargs,197            tokenizer_init_kwargs=self.tokenizer.init_kwargs,198            **kwargs,199        )200        format_output = output_kwargs["text_kwargs"].pop("format")201        num_image_tokens = output_kwargs["images_kwargs"].pop("num_image_tokens")202        box = output_kwargs["images_kwargs"].pop("box", [None])203        color = output_kwargs["images_kwargs"].pop("color", None)204        multi_page = output_kwargs["images_kwargs"].pop("multi_page")205 206        crop_to_patches = output_kwargs["images_kwargs"].get("crop_to_patches")207        images, text, box, color = self._make_list_of_inputs(images, text, box, color, multi_page)208        if multi_page:209            # save the number of pages per batch210            num_pages_per_batch = [len(image_group) for image_group in images]211            # flatten the list of images212            images = [image for image_group in images for image in image_group]213        else:214            num_pages_per_batch = [1 for _ in range(len(images))]215        # Load images as we need to know the image size216        images = load_images(images)217        image_sizes = [image.size for image in images]218        image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])219        num_patches_array = image_inputs.pop("num_patches")220        if text is None:221            text = []222            patch_indices = np.cumsum(num_pages_per_batch)223            for index, (num_pages, box_single, color_single) in enumerate(zip(num_pages_per_batch, box, color)):224                current_patch_index = patch_indices[index - 1] if index > 0 else 0225                num_patches = sum(num_patches_array[current_patch_index : current_patch_index + num_pages])226                if box_single[0] is not None:227                    box_single = preprocess_box_annotation(box_single, image_sizes[index])228                query = (229                    f"{f'[{color_single}] ' if color_single is not None else ''}"230                    f"{str(box_single) if box_single[0] is not None else ''} "231                    "OCR"232                    f"{' with format' if format_output else ''}"233                    f"{' across multi pages' if multi_page else ''}"234                    f"{' upon the patch reference' if crop_to_patches else ''}"235                    ": "236                )237                prompt = (238                    self.message_start_token239                    + self.system_query240                    + self.message_end_token241                    + self.message_start_token242                    + "user\n"243                    + self.img_start_token244                    + self.img_pad_token * num_image_tokens * num_patches245                    + self.img_end_token246                    + "\n"247                    + query248                    + self.message_end_token249                    + self.message_start_token250                    + "assistant\n"251                )252                text.append(prompt)253 254        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)255        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])256        self._check_special_mm_tokens(text, text_inputs, modalities=["image"])257 258        return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)259 260 261__all__ = ["GotOcr2Processor"]262