camenduru/kosmos-2-patch14-224

1# coding=utf-82# Copyright 2023 The HuggingFace Inc. team.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"""Processor class for KOSMOS-2."""16 17import copy18import math19import re20from typing import List, Optional, Tuple, Union21 22import numpy as np23 24from transformers.image_processing_utils import BatchFeature25from transformers.image_utils import ImageInput, is_batched26from transformers.processing_utils import ProcessorMixin27from transformers.tokenization_utils_base import PaddingStrategy, TextInput, TruncationStrategy28from transformers.utils import TensorType, is_tf_available, is_torch_available29 30 31if is_torch_available():32    import torch33 34if is_tf_available():35    import tensorflow as tf36 37 38BboxInput = Union[39    List[Tuple[int, int]],40    List[Tuple[float, float, float, float]],41    List[List[Tuple[int, int]]],42    List[List[Tuple[float, float, float]]],43]44 45 46class Kosmos2Processor(ProcessorMixin):47    r"""48    Constructs an KOSMOS-2 processor which wraps a CLIP image processor and a KOSMOS-2 tokenizer into a single49    processor.50 51    [`Kosmos2Processor`] offers all the functionalities of [`CLIPImageProcessor`] and [`Kosmos2TokenizerFast`]. See the52    docstring of [`~Kosmos2Processor.__call__`] and [`~Kosmos2Processor.decode`] for more information.53 54    Args:55        image_processor (`CLIPImageProcessor`):56            An instance of [`CLIPImageProcessor`]. The image processor is a required input.57        tokenizer (`Kosmos2TokenizerFast`):58            An instance of ['Kosmos2TokenizerFast`]. The tokenizer is a required input.59    """60    attributes = ["image_processor", "tokenizer"]61    # Better to use explicit classes if local code works62    # image_processor_class = "Kosmos2ImageProcessor"63    # tokenizer_class = ("Kosmos2Tokenizer", "Kosmos2TokenizerFast")64 65    # To make remote code work66    image_processor_class = "AutoImageProcessor"67    tokenizer_class = "AutoTokenizer"68 69    def __init__(self, image_processor, tokenizer):70        tokenizer.return_token_type_ids = False71        super().__init__(image_processor, tokenizer)72        self.current_processor = self.image_processor73 74    def __call__(75        self,76        images: ImageInput = None,77        text: Union[TextInput, List[TextInput]] = None,78        bboxes: BboxInput = None,79        num_image_tokens: Optional[int] = 64,80        first_image_token_id: Optional[int] = None,81        add_special_tokens: bool = True,82        padding: Union[bool, str, PaddingStrategy] = False,83        truncation: Union[bool, str, TruncationStrategy] = None,84        max_length: Optional[int] = None,85        stride: int = 0,86        pad_to_multiple_of: Optional[int] = None,87        return_attention_mask: Optional[bool] = None,88        return_overflowing_tokens: bool = False,89        return_special_tokens_mask: bool = False,90        return_offsets_mapping: bool = False,91        return_token_type_ids: bool = False,92        return_length: bool = False,93        verbose: bool = True,94        return_tensors: Optional[Union[str, TensorType]] = None,95        **kwargs,96    ) -> BatchFeature:97        """98        This method uses [`CLIPImageProcessor.__call__`] method to prepare image(s) for the model, and99        [`Kosmos2TokenizerFast.__call__`] to prepare text for the model.100 101        Please refer to the docstring of the above two methods for more information.102        """103        if text is None:104            raise ValueError("You have to specify at least `text`.")105 106        text = self.preprocess_text(text, images, bboxes, num_image_tokens=num_image_tokens)107 108        encoding = BatchFeature()109 110        text_encoding = self.tokenizer(111            text=text,112            add_special_tokens=add_special_tokens,113            padding=padding,114            truncation=truncation,115            max_length=max_length,116            stride=stride,117            pad_to_multiple_of=pad_to_multiple_of,118            return_attention_mask=return_attention_mask,119            return_overflowing_tokens=return_overflowing_tokens,120            return_special_tokens_mask=return_special_tokens_mask,121            return_offsets_mapping=return_offsets_mapping,122            return_token_type_ids=return_token_type_ids,123            return_length=return_length,124            verbose=verbose,125            return_tensors=return_tensors,126            **kwargs,127        )128        encoding.update(text_encoding)129 130        if images is not None:131            image_encoding = self.image_processor(images, return_tensors=return_tensors)132            encoding.update(image_encoding)133 134            # Use the id of the first token after <unk>135            if first_image_token_id is None:136                first_image_token_id = self.tokenizer.unk_token_id + 1137 138            # To see if we need one more `0` (for `<s>`) at the beginning of `img_attn_mask`.139            with_bos = add_special_tokens140 141            # The first (actual) `<image>` token is always at the 1st or 2nd place (after `<s>` if any). Here we look142            # for the second `<image>` token (which indicate the first image token).143            start_index = int(with_bos) + 1144 145            if return_tensors:146                # change the ids for the fake `<image>` tokens in `input_ids`147                input_ids = np.array(encoding["input_ids"])148                input_ids[:, start_index : (start_index + num_image_tokens)] = np.arange(149                    first_image_token_id, first_image_token_id + num_image_tokens150                )151 152                batch_size, seq_len = input_ids.shape[:2]153                img_attn_mask = []154                if with_bos:155                    # for `<s>`156                    img_attn_mask.append(np.zeros(shape=(batch_size, 1), dtype=np.int64))157                # for `<image>` (the real one)158                img_attn_mask.append(np.zeros(shape=(batch_size, 1), dtype=np.int64))159                # for image tokens160                img_attn_mask.append(np.ones(shape=(batch_size, 64), dtype=np.int64))161                # for `</image>`162                img_attn_mask.append(np.zeros(shape=(batch_size, 1), dtype=np.int64))163                # trailing part (which are not related to the image)164                seq_len -= int(with_bos) + 1 + num_image_tokens + 1165                img_attn_mask.append(np.zeros(shape=(batch_size, seq_len), dtype=np.int64))166 167                # concatenate along the sequence dimension168                img_attn_mask = np.concatenate(img_attn_mask, axis=1)169 170                # to the target tensor type171                if return_tensors == "pt":172                    input_ids = torch.from_numpy(input_ids)173                    img_attn_mask = torch.from_numpy(img_attn_mask)174                elif return_tensors == "tf":175                    input_ids = tf.convert_to_tensor(input_ids)176                    img_attn_mask = tf.convert_to_tensor(img_attn_mask)177 178                encoding["input_ids"] = input_ids179                encoding["img_attn_mask"] = img_attn_mask180 181            else:182                # Add `img_attn_mask`: the leading and trailing `0` are for `boi` and `eoi` tokens. The `1` indicates183                # the places of image tokens.184                image_token_ids = list(range(first_image_token_id, first_image_token_id + num_image_tokens))185                base_img_attn_mask = [0] + [1] * num_image_tokens + [0]186 187                # loop over `encoding["input_ids"]`188                input_ids = []189                img_attn_mask = []190                all_input_ids = encoding["input_ids"]191                # not batched -> (changed to) batch of size 1192                if isinstance(text, str):193                    all_input_ids = [all_input_ids]194                for text_ids in all_input_ids:195                    # change the ids for the fake `<image>` tokens in `input_ids`196                    text_ids = text_ids[:start_index] + image_token_ids + text_ids[start_index + num_image_tokens :]197                    input_ids.append(text_ids)198 199                    mask = copy.copy(base_img_attn_mask)200                    if with_bos:201                        # for `<s>`202                        mask = [0] + mask203                    # trailing part (which are not related to the image)204                    mask += [0] * (len(text_ids) - len(mask))205                    img_attn_mask.append(mask)206 207                # un-batch if necessary208                if isinstance(text, str):209                    input_ids = input_ids[0]210                    img_attn_mask = img_attn_mask[0]211 212                encoding["input_ids"] = input_ids213                encoding["img_attn_mask"] = img_attn_mask214 215        return encoding216 217    def preprocess_text(218        self,219        texts: Union[TextInput, List[TextInput]],220        images: ImageInput = None,221        bboxes: BboxInput = None,222        num_image_tokens: Optional[int] = 64,223    ) -> Union[str, List[str]]:224        """Add image and bounding box information to `texts` as image and patch index tokens.225 226        Args:227            texts (`Union[TextInput, List[TextInput]]`): The texts to be processed.228            images (`ImageInput`, *optional*): The images associated to `texts`.229            bboxes (`Union[List[Tuple[int]], List[Tuple[float]], List[List[Tuple[int]]], List[List[Tuple[float]]]]`, *optional*): The bounding bboxes associated to `texts`.230            num_image_tokens (`int`, *optional*, defaults to 64): The number of image tokens (used as latent queries). This should corresponds to the `latent_query_num` attribute in `Kosmos2Config`.231 232        Returns:233            `Union[TextInput, List[TextInput]]`: The processed texts with image and patch index tokens.234        """235        # These are fake `<image>` tokens enclosed between (the actual) `<image>` token and `</image>`.236        img_tokens = ["<image>"] * num_image_tokens237        img_info = " ".join(["<image>"] + img_tokens + ["</image>"])238 239        def check_bboxes_for_single_text(bboxes):240            """241            Check `bboxes` for a single text example. It could be242                - `None`: no bounding box associated to a text.243                - A list with each element being the bounding boxes associated to one `<phrase> ... </phrase>` pair244                  found in a text. This could be:245                      - `None`: no bounding box associated to a `<phrase> ... </phrase>` pair.246                      - A tuple of 2 integers: A single bounding box specified by patch indices.247                      - A tuple of 4 float point number: A single bounding box specified by (normalized) coordinates.248                      - A list containing the above 2 tuple types: Multiple bounding boxes for a249                       `<phrase> ... </phrase>` pair.250            """251            if bboxes is None:252                return253            elif not isinstance(bboxes, list):254                raise ValueError("`bboxes` (for a single text example) should be `None` or a list.")255 256            # `bbox` is the bounding boxes for a single <phrase> </phrase> pair257            for bbox in bboxes:258                if bbox is None:259                    continue260                elif not isinstance(bbox, list):261                    bbox = [bbox]262                for elt in bbox:263                    if not isinstance(elt, tuple) or not (264                        (len(elt) == 2 and all(isinstance(x, int) for x in elt))265                        or (len(elt) == 4 and all(isinstance(x, float) for x in elt))266                    ):267                        raise ValueError(268                            "Each element in `bboxes` (for a single text example) should be `None`, a tuple containing "269                            "2 integers or 4 float point numbers, or a list containing such tuples. Also "270                            "make sure the arguments `texts` and `bboxes` passed to `preprocess_text` are both in "271                            "batches or both for a single example."272                        )273 274        def preprocess_single(text, image, bboxes):275            text = text.strip()276            if image is not None:277                # Add `<image> ... (fake) image tokens ... </image>`278                text = f"{img_info} {text}"279 280            # Add `<object> <patch_idx_xxxx> <patch_idx_yyy> </object>` after `<phrase> phrase text </phrase>`281            text = self._insert_patch_index_tokens(text, bboxes)282            text = self._add_remove_spaces_around_tag_tokens(text)283 284            return text285 286        # make batch to simplify processing logic287        batched = True288        if isinstance(texts, str):289            batched = False290            texts = [texts]291 292        if images is None:293            images = [None] * len(texts)294        elif not is_batched(images):295            images = [images]296        if len(texts) != len(images):297            raise ValueError(298                f"The number of examples in `texts` and `images` should be the same. Got {len(texts)} v.s. {len(images)} instead."299            )300 301        if not batched:302            check_bboxes_for_single_text(bboxes)303            bboxes = [bboxes]304        elif bboxes is not None:305            if not isinstance(bboxes, list):306                raise ValueError("`bboxes` should be `None` or a list (as a batch) when `texts` is passed as a batch.")307            for x in bboxes:308                check_bboxes_for_single_text(x)309        else:310            bboxes = [None] * len(texts)311 312        if len(bboxes) != len(texts):313            raise ValueError(314                f"The number of examples in `texts` and `bboxes` should be the same. Got {len(texts)} v.s. {len(bboxes)} instead."315            )316 317        result = [preprocess_single(text, image, bbox) for text, image, bbox in zip(texts, images, bboxes)]318        # un-batch if necessary319        if not batched:320            result = result[0]321 322        return result323 324    # Copied from transformers.models.blip.processing_blip.BlipProcessor.batch_decode with BertTokenizerFast->PreTrainedTokenizer325    def batch_decode(self, *args, **kwargs):326        """327        This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please328        refer to the docstring of this method for more information.329        """330        return self.tokenizer.batch_decode(*args, **kwargs)331 332    # Copied from transformers.models.blip.processing_blip.BlipProcessor.decode with BertTokenizerFast->PreTrainedTokenizer333    def decode(self, *args, **kwargs):334        """335        This method forwards all its arguments to PreTrainedTokenizer's [`~PreTrainedTokenizer.decode`]. Please refer336        to the docstring of this method for more information.337        """338        return self.tokenizer.decode(*args, **kwargs)339 340    def post_process_generation(self, text, cleanup_and_extract=True):341 342        caption = text.split("</image>")[-1]343        if cleanup_and_extract:344            return clean_text_and_extract_entities_with_bboxes(caption)345        return caption346 347    @property348    # Copied from transformers.models.blip.processing_blip.BlipProcessor.model_input_names349    def model_input_names(self):350        tokenizer_input_names = self.tokenizer.model_input_names351        image_processor_input_names = self.image_processor.model_input_names352        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))353 354    def _insert_patch_index_tokens(self, text: str, bboxes: Union[List[Tuple[int]], List[Tuple[float]]]) -> str:355        if bboxes is None or len(bboxes) == 0:356            return text357 358        matched_phrases = list(re.finditer(r"<phrase>.+?</phrase>", string=text))359        if len(matched_phrases) != len(bboxes):360            raise ValueError(361                f"The number of elements in `bboxes` should be the same as the number of `<phrase> ... </phrase>` pairs in `text`. Got {len(matched_phrases)} v.s. {len(bboxes)} instead."362            )363 364        # insert object's patch index tokens365        # the found `<phrase> ... </phrase>` pairs.366        curr_pos = 0367        buffer = []368        for matched, bbox in zip(matched_phrases, bboxes):369            _, end = matched.span()370            buffer.append(text[curr_pos:end])371            curr_pos = end372            # A phrase without bbox373            if bbox is None:374                continue375            # A phrase with a single bbox376            if isinstance(bbox, tuple):377                bbox = [bbox]378            patch_index_strings = []379            # A phrase could have multiple bboxes380            for box in bbox:381                patch_index_1, patch_index_2 = self._convert_bbox_to_patch_index_tokens(box)382                patch_index_strings.append(f"{patch_index_1} {patch_index_2}")383            position_str = " </delimiter_of_multi_objects/> ".join(patch_index_strings)384            buffer.append(f"<object> {position_str} </object>")385        # remaining386        if curr_pos < len(text):387            buffer.append(text[curr_pos:])388 389        text = "".join(buffer)390        return text391 392    def _convert_bbox_to_patch_index_tokens(393        self, bbox: Union[Tuple[int, int], Tuple[float, float, float, float]]394    ) -> Tuple[str, str]:395        # already computed patch indices396        if len(bbox) == 2:397            idx_1, idx_2 = bbox398        # bbox specified with (normalized) coordinates399        else:400            # use `self.tokenizer` to get `num_patches_per_side`401            num_patches_per_side = int(math.sqrt(self.tokenizer.num_patch_index_tokens))402            idx_1, idx_2 = coordinate_to_patch_index(bbox, num_patches_per_side)403 404        token_1 = f"<patch_index_{str(idx_1).zfill(4)}>"405        token_2 = f"<patch_index_{str(idx_2).zfill(4)}>"406 407        return token_1, token_2408 409    def _add_remove_spaces_around_tag_tokens(self, text):410        """411        Remove spaces before tag tokens (e.g. `<x>`). Also ensure a space after a tag token, if it is not followed by412        another tag token (this is not technically necessary, but good for a standard/consistent format). This avoids413        the inconsistency of tokenization results between kosmos-2 slow and fast tokenizers.414        """415 416        tag_tokens = set(417            self.tokenizer.tag_tokens418            + [f"<patch_index_{str(x).zfill(4)}>" for x in range(self.tokenizer.num_patch_index_tokens)]419        )420        pattern = "|".join(tag_tokens)421        splits = re.split(rf"({pattern})", text)422        # Don't keep the leading and trailing space if any423        splits = [split for idx, split in enumerate(splits) if not (idx in [0, len(splits) - 1] and split == "")]424 425        output = ""426        prev_str_in_targets = False427        for split in splits:428            if split in tag_tokens:429                prev_str_in_targets = True430                output = output.rstrip() + split431            else:432                # we don't need to ensure a space before a normal token that is after a tag token. But having it and433                # keeps a standard format is good anyway.434                if prev_str_in_targets and not split.startswith(" "):435                    output += " " + split436                else:437                    output += split438                prev_str_in_targets = False439 440        return output441 442 443def coordinate_to_patch_index(bbox: Tuple[float, float, float, float], num_patches_per_side: int) -> Tuple[int, int]:444    """Convert a bounding box to a pair of patch indices.445 446    Args:447        bbox (`Tuple[float, float, float, float]`):448            The 4 coordinates of the bounding box, with the format being (x1, y1, x2, y2) specifying the upper-left449            and lower-right corners of the box. It should have x2 > x1 and  y1 > y2.450        num_patches_per_side (`int`): the number of patches along each side.451 452    Returns:453        `Tuple[int, int]`: A pair of patch indices.454    """455    (x1, y1, x2, y2) = bbox456 457    ul_x = math.floor(x1 * num_patches_per_side)458    ul_y = math.floor(y1 * num_patches_per_side)459 460    lr_x = math.ceil(x2 * num_patches_per_side - 1)461    lr_y = math.ceil(y2 * num_patches_per_side - 1)462 463    ul_idx = ul_y * num_patches_per_side + ul_x464    lr_idx = lr_y * num_patches_per_side + lr_x465 466    return ul_idx, lr_idx467 468 469# copied from https://github.com/microsoft/unilm/blob/97e4923e97d3ee10b57e97013556e3fd0d207a9b/kosmos-2/demo/decode_string.py#L35C1-L75C38470# (with format modifications)471def patch_index_to_coordinate(ul_idx: int, lr_idx: int, num_patches_per_side: int):472    """473    Given a grid of length `num_patches_per_side` and the indices of the upper-left and lower-right corners of a474    bounding box, returns the normalized coordinates of the bounding box, in the form (x1, y1, x2, y2).475 476    Args:477        ul_idx (`int`): the index of the grid cell that corresponds to the upper-left corner of the bounding box.478        lr_idx (`int`): the index of the grid cell that corresponds to the lower-right corner of the bounding box.479        num_patches_per_side (`int`): the number of patches along each side.480 481    Returns:482        `Tuple[float]`: the normalized coordinates of the bounding box, in the form (x1, y1, x2, y2).483    """484    # Compute the size of each cell in the grid485    cell_size = 1.0 / num_patches_per_side486 487    # Compute the x and y indices of the upper-left and lower-right corners of the bounding box488    ul_x = ul_idx % num_patches_per_side489    ul_y = ul_idx // num_patches_per_side490 491    lr_x = lr_idx % num_patches_per_side492    lr_y = lr_idx // num_patches_per_side493 494    # Compute the normalized coordinates of the bounding box495    if ul_idx == lr_idx:496        x1 = ul_x * cell_size497        y1 = ul_y * cell_size498        x2 = lr_x * cell_size + cell_size499        y2 = lr_y * cell_size + cell_size500    elif ul_x == lr_x or ul_y == lr_y:501        x1 = ul_x * cell_size502        y1 = ul_y * cell_size503        x2 = lr_x * cell_size + cell_size504        y2 = lr_y * cell_size + cell_size505    else:506        x1 = ul_x * cell_size + cell_size / 2507        y1 = ul_y * cell_size + cell_size / 2508        x2 = lr_x * cell_size + cell_size / 2509        y2 = lr_y * cell_size + cell_size / 2510 511    return x1, y1, x2, y2512 513 514# copied from https://github.com/microsoft/unilm/blob/97e4923e97d3ee10b57e97013556e3fd0d207a9b/kosmos-2/demo/decode_string.py#L4-L33515# (with format modifications)516def extract_entities_with_patch_indices(text):517    # The regular expression pattern for matching the required formats518    pattern = r'(?:(<phrase>([^<]+)</phrase>))?<object>((?:<patch_index_\d+><patch_index_\d+></delimiter_of_multi_objects/>)*<patch_index_\d+><patch_index_\d+>)</object>'519 520    # Find all matches in the given string521    matches = re.finditer(pattern, text)522 523    # Initialize an empty list to store the valid patch_index combinations524    entities_with_patch_indices = []525 526    for match in matches:527        # span of a `phrase` that is between <phrase> and </phrase>528        span = match.span(2)529        phrase_tag, phrase, match_content = match.groups()530        if not phrase_tag:531            phrase = None532            # We take the starting position of `<object>`533            span = (match.span(0)[0], match.span(0)[0])534 535        # Split the match_content by the delimiter to get individual patch_index pairs536        patch_index_pairs = match_content.split('</delimiter_of_multi_objects/>')537 538        entity_bboxes = []539        for pair in patch_index_pairs:540            # Extract the xxxx and yyyy values from the patch_index pair541            x = re.search(r'<patch_index_(\d+)>', pair)542            y = re.search(r'<patch_index_(\d+)>', pair[1:])543 544            if x and y:545                if phrase:546                    entity_bboxes.append((int(x.group(1)), int(y.group(1))))547                else:548                    entity_bboxes.append((int(x.group(1)), int(y.group(1))))549 550        if phrase:551            entities_with_patch_indices.append((phrase, span, entity_bboxes))552        else:553            for bbox in entity_bboxes:554                # fake entity name555                entity = f"<patch_index_{bbox[0]}><patch_index_{bbox[1]}>"556                entities_with_patch_indices.append((entity, span, [bbox]))557 558    return entities_with_patch_indices559 560 561# TODO: Be careful562def remove_special_fields(text):563    return re.sub('<.*?>', '', text)564 565 566def adjust_entity_positions(entity, text):567 568    entity_name, (start, end) = entity569    adjusted_start = len(remove_special_fields(text[:start]))570    adjusted_end = len(remove_special_fields(text[:end]))571    adjusted_entity = (entity_name, (adjusted_start, adjusted_end))572    return adjusted_entity573 574 575# copied from https://github.com/microsoft/unilm/blob/97e4923e97d3ee10b57e97013556e3fd0d207a9b/kosmos-2/demo/decode_string.py#L77-L87576# (with format modifications)577def clean_text_and_extract_entities_with_bboxes(text, num_patches_per_side=32):578 579    processed_text = remove_special_fields(text)580 581    entities_with_patch_indices = extract_entities_with_patch_indices(text)582    entities = []583    for item in entities_with_patch_indices:584        entity, bboxes = item[0:2], item[2]585        adjusted_entity = adjust_entity_positions(entity, text)586        bboxes_in_coords = list(map(lambda bbox: patch_index_to_coordinate(bbox[0], bbox[1], num_patches_per_side), bboxes))587 588        entities.append(adjusted_entity + (bboxes_in_coords,))589 590    def cleanup_spaces(text, entities):591        new_text = text.strip()592        leading_spaces = len(text) - len(text.lstrip())593 594        new_entities = []595        for entity_name, (start, end), bboxes in entities:596 597            entity_name_leading_spaces = len(entity_name) - len(entity_name.lstrip())598            entity_name_trailing_spaces = len(entity_name) - len(entity_name.rstrip())599 600            start = start - leading_spaces + entity_name_leading_spaces601            end = end - leading_spaces - entity_name_trailing_spaces602            entity_name = entity_name.strip()603 604            new_entities.append((entity_name, (start, end), bboxes))605 606        return new_text, new_entities607 608    return cleanup_spaces(processed_text, entities)609