processing_florence2.py

# coding=utf-8
# Copyright 2024 Microsoft and 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.
#
# This file has been modified from its original version by Julian Spravil, 2025.

"""
Processor class for Florence-2.
"""

import logging
import re
from typing import List, Optional, Union

import numpy as np
import torch
from transformers import BartTokenizer, BartTokenizerFast, T5Tokenizer, T5TokenizerFast
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_utils import ImageInput, is_valid_image
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import (
    PaddingStrategy,
    PreTokenizedInput,
    TextInput,
    TruncationStrategy,
)
from transformers.utils import TensorType

logger = logging.getLogger(__name__)


# Copied from transformers.models.idefics2.processing_idefics2.is_url
def is_url(val) -> bool:
    return isinstance(val, str) and val.startswith("http")


# Copied from transformers.models.idefics2.processing_idefics2.is_image_or_image_url
def is_image_or_image_url(elem):
    return is_url(elem) or is_valid_image(elem)


def _is_str_or_image(elem):
    return isinstance(elem, (str)) or is_image_or_image_url(elem)


class Florence2Processor(ProcessorMixin):
    r"""
    Constructs a Florence2 processor which wraps a Florence2 image processor and a Florence2 tokenizer into a single processor.

    [`Florence2Processor`] offers all the functionalities of [`CLIPImageProcessor`] and [`BartTokenizerFast`]. See the
    [`~Florence2Processor.__call__`] and [`~Florence2Processor.decode`] for more information.

    Args:
        image_processor ([`CLIPImageProcessor`], *optional*):
            The image processor is a required input.
        tokenizer ([`BartTokenizerFast`], *optional*):
            The tokenizer is a required input.
    """

    attributes = ["image_processor", "tokenizer"]
    image_processor_class = "CLIPImageProcessor"
    tokenizer_class = ("BartTokenizer", "BartTokenizerFast")
    valid_kwargs = ["new_tokenizer", "use_encoder_tokenizer", "insert_lang_token"]

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        new_tokenizer=None,
        use_encoder_tokenizer=False,
        insert_lang_token=False,
    ):
        if image_processor is None:
            raise ValueError("You need to specify an `image_processor`.")
        if tokenizer is None:
            raise ValueError("You need to specify a `tokenizer`.")
        if not hasattr(image_processor, "image_seq_length"):
            raise ValueError("Image processor is missing an `image_seq_length` attribute.")

        self.encoder_tokenizer = None
        if use_encoder_tokenizer:
            assert (
                new_tokenizer is not None
            ), "Encoder tokenizer can only be used if a new tokenizer for the decoder is provided"
            print(f"Use {type(tokenizer).__name__} as encoder tokenizer ...")
            self.encoder_tokenizer = tokenizer

        if new_tokenizer is not None:
            print(f"Replacing decoder tokenizer {type(tokenizer).__name__} with {type(new_tokenizer).__name__} ...")
            self.tokenizer_class = new_tokenizer.__class__.__name__
            new_tokenizer.model_max_length = tokenizer.model_max_length
            tokenizer = new_tokenizer

        self.image_seq_length = image_processor.image_seq_length

        if self.encoder_tokenizer is not None:
            tokens_to_add = {
                "additional_special_tokens": self.encoder_tokenizer.additional_special_tokens
                + ["<od>", "</od>", "<ocr>", "</ocr>"]
                + [f"<loc_{x}>" for x in range(1000)]
                + [
                    "<cap>",
                    "</cap>",
                    "<ncap>",
                    "</ncap>",
                    "<dcap>",
                    "</dcap>",
                    "<grounding>",
                    "</grounding>",
                    "<seg>",
                    "</seg>",
                    "<sep>",
                    "<region_cap>",
                    "</region_cap>",
                    "<region_to_desciption>",
                    "</region_to_desciption>",
                    "<proposal>",
                    "</proposal>",
                    "<poly>",
                    "</poly>",
                    "<and>",
                ]
            }

            self.encoder_tokenizer.add_special_tokens(tokens_to_add)
        else:
            tokens_to_add = {
                "additional_special_tokens": tokenizer.additional_special_tokens
                + ["<od>", "</od>", "<ocr>", "</ocr>"]
                + [f"<loc_{x}>" for x in range(1000)]
                + [
                    "<cap>",
                    "</cap>",
                    "<ncap>",
                    "</ncap>",
                    "<dcap>",
                    "</dcap>",
                    "<grounding>",
                    "</grounding>",
                    "<seg>",
                    "</seg>",
                    "<sep>",
                    "<region_cap>",
                    "</region_cap>",
                    "<region_to_desciption>",
                    "</region_to_desciption>",
                    "<proposal>",
                    "</proposal>",
                    "<poly>",
                    "</poly>",
                    "<and>",
                ]
            }
            tokenizer.add_special_tokens(tokens_to_add)

        if insert_lang_token:
            tokenizer.add_special_tokens(
                {
                    "additional_special_tokens": tokenizer.additional_special_tokens
                    + ["<LANG_EN>", "<LANG_DE>", "<LANG_ES>", "<LANG_ZH>", "<LANG_RU>"]
                }
            )

        self.tasks_answer_post_processing_type = {
            "<OCR>": "pure_text",
            "<OCR_WITH_REGION>": "ocr",
            "<CAPTION>": "pure_text",
            "<DETAILED_CAPTION>": "pure_text",
            "<MORE_DETAILED_CAPTION>": "pure_text",
            "<OD>": "description_with_bboxes",
            "<DENSE_REGION_CAPTION>": "description_with_bboxes",
            "<CAPTION_TO_PHRASE_GROUNDING>": "phrase_grounding",
            "<REFERRING_EXPRESSION_SEGMENTATION>": "polygons",
            "<REGION_TO_SEGMENTATION>": "polygons",
            "<OPEN_VOCABULARY_DETECTION>": "description_with_bboxes_or_polygons",
            "<REGION_TO_CATEGORY>": "pure_text",
            "<REGION_TO_DESCRIPTION>": "pure_text",
            "<REGION_TO_OCR>": "pure_text",
            "<REGION_PROPOSAL>": "bboxes",
            "<TRANSLATE_TO_GERMAN>": "pure_text",
        }

        self.task_prompts_without_inputs = {
            "<OCR>": "What is the text in the image?",
            "<OCR_WITH_REGION>": "What is the text in the image, with regions?",
            "<CAPTION>": "What does the image describe?",
            "<DETAILED_CAPTION>": "Describe in detail what is shown in the image.",
            "<MORE_DETAILED_CAPTION>": "Describe with a paragraph what is shown in the image.",
            "<OD>": "Locate the objects with category name in the image.",
            "<DENSE_REGION_CAPTION>": "Locate the objects in the image, with their descriptions.",
            "<REGION_PROPOSAL>": "Locate the region proposals in the image.",
        }

        self.task_prompts_with_input = {
            "<CAPTION_TO_PHRASE_GROUNDING>": "Locate the phrases in the caption: {input}",
            "<REFERRING_EXPRESSION_SEGMENTATION>": "Locate {input} in the image with mask",
            "<REGION_TO_SEGMENTATION>": "What is the polygon mask of region {input}",
            "<OPEN_VOCABULARY_DETECTION>": "Locate {input} in the image.",
            "<REGION_TO_CATEGORY>": "What is the region {input}?",
            "<REGION_TO_DESCRIPTION>": "What does the region {input} describe?",
            "<REGION_TO_OCR>": "What text is in the region {input}?",
            "<TRANSLATE>": "Translate the following text: {input}",
        }

        self.task_modifiers = {
            "<LANG_EN>": "The target language is English.",
            "<LANG_DE>": "The target language is German.",
            "<LANG_FR>": "The target language is French.",
            "<LANG_ES>": "The target language is Spanish.",
            "<LANG_ZH>": "The target language is simplified Chinese.",
            "<LANG_RU>": "The target language is Russian.",
        }

        self.post_processor = Florence2PostProcesser(tokenizer=tokenizer)

        super().__init__(image_processor, tokenizer)

    @classmethod
    def from_args_and_dict(cls, args, processor_dict, **kwargs):
        # Make sure that the kwarg for new tokenizer is passed correctly to the init function
        if "new_tokenizer" in kwargs:
            processor_dict["new_tokenizer"] = kwargs["new_tokenizer"]
        if "use_encoder_tokenizer" in kwargs:
            processor_dict["use_encoder_tokenizer"] = kwargs["use_encoder_tokenizer"]
        if "insert_lang_token" in kwargs:
            processor_dict["insert_lang_token"] = kwargs["insert_lang_token"]
        return super().from_args_and_dict(args, processor_dict, **kwargs)

    def _construct_prompts(self, text):
        # replace the task tokens with the task prompts if task token is in the text
        prompts = []
        for _text in text:
            segments = []
            for task_token, task_prompt in self.task_modifiers.items():
                if task_token in _text:
                    _text = _text.replace(task_token, "")
                    segments.append(task_prompt)
            assert len(segments) <= 1, "Only one modifier is allowed."

            # 1. fixed task prompts without additional inputs
            for task_token, task_prompt in self.task_prompts_without_inputs.items():
                if task_token in _text:
                    assert _text == task_token, f"Task token {task_token} should be the only token in the text."
                    _text = task_prompt
                    break
            # 2. task prompts with additional inputs
            for task_token, task_prompt in self.task_prompts_with_input.items():
                if task_token in _text:
                    _text = task_prompt.format(input=_text.replace(task_token, ""))
                    break
            segments.append(_text)
            prompts.append(" ".join(segments))
        return prompts

    def __call__(
        self,
        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
        images: ImageInput = None,
        tokenize_newline_separately: bool = True,
        padding: Union[bool, str, PaddingStrategy] = False,
        truncation: Union[bool, str, TruncationStrategy] = None,
        max_length=None,
        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
        do_resize: bool = None,
        do_normalize: bool = None,
        image_mean: Optional[Union[float, List[float]]] = None,
        image_std: Optional[Union[float, List[float]]] = None,
        data_format: Optional["ChannelDimension"] = "channels_first",  # noqa: F821
        input_data_format: Optional[
            Union[str, "ChannelDimension"]  # noqa: F821
        ] = None,
        resample: "PILImageResampling" = None,  # noqa: F821
        do_convert_rgb: bool = None,
        do_thumbnail: bool = None,
        do_align_long_axis: bool = None,
        do_rescale: bool = None,
        return_attention_mask: bool = None,
    ) -> BatchFeature:
        """
        Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
        and `kwargs` arguments to BartTokenizerFast's [`~BartTokenizerFast.__call__`] if `text` is not `None` to encode
        the text. To prepare the image(s), this method forwards the `images` and `kwrags` arguments to
        CLIPImageProcessor's [`~CLIPImageProcessor.__call__`] if `images` is not `None`. Please refer to the doctsring
        of the above two methods for more information.

        Args:
            text (`str`, `List[str]`, `List[List[str]]`):
                The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
                (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
                `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
                tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
                number of channels, H and W are image height and width.
            tokenize_newline_separately (`bool`, defaults to `True`):
                Adds a separately tokenized '\n' at the end of the prompt.
            padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `False`):
                Select a strategy to pad the returned sequences (according to the model's padding side and padding
                index) among:
                - `True` or `'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
                  sequence if provided).
                - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
                  acceptable input length for the model if that argument is not provided.
                - `False` or `'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of different
                  lengths).
            max_length (`int`, *optional*):
                Maximum length of the returned list and optionally padding length (see above).
            truncation (`bool`, *optional*):
                Activates truncation to cut input sequences longer than `max_length` to `max_length`.
            return_tensors (`str` or [`~utils.TensorType`], *optional*):
                If set, will return tensors of a particular framework. Acceptable values are:

                - `'tf'`: Return TensorFlow `tf.constant` objects.
                - `'pt'`: Return PyTorch `torch.Tensor` objects.
                - `'np'`: Return NumPy `np.ndarray` objects.
                - `'jax'`: Return JAX `jnp.ndarray` objects.

        Returns:
            [`BatchFeature`]: A [`BatchFeature`] with the following fields:

            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`. If `suffix`
              is provided, the `input_ids` will also contain the suffix input ids.
            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
              `None`).
            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
            - **labels** -- Labels compatible with training if `suffix` is not None
        """

        return_token_type_ids = False

        if images is None:
            raise ValueError("`images` are expected as arguments to a `Florence2Processor` instance.")
        if text is None:
            logger.warning_once("You are using Florence-2 without a text prompt.")
            text = ""

        if isinstance(text, List) and isinstance(images, List):
            if len(images) < len(text):
                raise ValueError(
                    f"Received {len(images)} images for {len(text)} prompts. Each prompt should be associated with an image."
                )
        if _is_str_or_image(text):
            text = [text]
        elif isinstance(text, list) and _is_str_or_image(text[0]):
            pass

        pixel_values = self.image_processor(
            images,
            do_resize=do_resize,
            do_normalize=do_normalize,
            return_tensors=return_tensors,
            image_mean=image_mean,
            image_std=image_std,
            input_data_format=input_data_format,
            data_format=data_format,
            resample=resample,
            do_convert_rgb=do_convert_rgb,
        )["pixel_values"]

        if max_length is not None:
            max_length -= self.image_seq_length  # max_length has to account for the image tokens

        text = self._construct_prompts(text)

        tokenizer_call = self.tokenizer
        if self.encoder_tokenizer is not None:
            tokenizer_call = self.encoder_tokenizer
        inputs = tokenizer_call(
            text,
            return_tensors=return_tensors,
            padding=padding,
            max_length=max_length,
            truncation=truncation,
            return_token_type_ids=return_token_type_ids,
            return_attention_mask=return_attention_mask,
        )

        return_data = {**inputs, "pixel_values": pixel_values}

        if return_token_type_ids:
            if self.encoder_tokenizer is not None:
                raise NotImplementedError("return_token_type_ids is not implemented with encoder tokenizer")
            labels = inputs["input_ids"].masked_fill(inputs["token_type_ids"] == 0, -100)
            return_data.update({"labels": labels})
        return BatchFeature(data=return_data)

    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.batch_decode with CLIP->Florence2
    def batch_decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to BartTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
        refer to the docstring of this method for more information.
        """
        return self.tokenizer.batch_decode(*args, **kwargs)

    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.decode with CLIP->Florence2
    def decode(self, *args, **kwargs):
        """
        This method forwards all its arguments to BartTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
        the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, **kwargs)

    @property
    # Copied from transformers.models.clip.processing_clip.CLIPProcessor.model_input_names with CLIP->Florence2
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))

    def post_process_generation(self, text, task, image_size):
        """
        Post-process the output of the model to each of the task outputs.

        Args:
            text (`str`): The text to post-process.
            task (`str`): The task to post-process the text for.
            image_size (`Tuple[int, int]`): The size of the image. height x width.
        """

        task_answer_post_processing_type = self.tasks_answer_post_processing_type.get(task, "pure_text")
        task_answer = self.post_processor(
            text=text,
            image_size=image_size,
            parse_tasks=task_answer_post_processing_type,
        )[task_answer_post_processing_type]

        if task_answer_post_processing_type == "pure_text":
            final_answer = task_answer
            # remove the special tokens
            final_answer = final_answer.replace("<eos>", "").replace("<bos>", "")
            final_answer = final_answer.replace("<s>", "").replace("</s>", "")
        elif task_answer_post_processing_type in ["od", "description_with_bboxes", "bboxes"]:
            od_instances = task_answer
            bboxes_od = [_od_instance["bbox"] for _od_instance in od_instances]
            labels_od = [str(_od_instance["cat_name"]) for _od_instance in od_instances]
            final_answer = {"bboxes": bboxes_od, "labels": labels_od}
        elif task_answer_post_processing_type in ["ocr"]:
            bboxes = [_od_instance["quad_box"] for _od_instance in task_answer]
            labels = [str(_od_instance["text"]) for _od_instance in task_answer]
            final_answer = {"quad_boxes": bboxes, "labels": labels}
        elif task_answer_post_processing_type in ["phrase_grounding"]:
            bboxes = []
            labels = []
            for _grounded_phrase in task_answer:
                for _bbox in _grounded_phrase["bbox"]:
                    bboxes.append(_bbox)
                    labels.append(_grounded_phrase["cat_name"])
            final_answer = {"bboxes": bboxes, "labels": labels}
        elif task_answer_post_processing_type in ["description_with_polygons", "polygons"]:
            labels = []
            polygons = []
            for result in task_answer:
                label = result["cat_name"]
                _polygons = result["polygons"]
                labels.append(label)
                polygons.append(_polygons)
            final_answer = {"polygons": polygons, "labels": labels}
        elif task_answer_post_processing_type in ["description_with_bboxes_or_polygons"]:
            bboxes = []
            bboxes_labels = []
            polygons = []
            polygons_labels = []
            for result in task_answer:
                label = result["cat_name"]
                if "polygons" in result:
                    _polygons = result["polygons"]
                    polygons.append(_polygons)
                    polygons_labels.append(label)
                else:
                    _bbox = result["bbox"]
                    bboxes.append(_bbox)
                    bboxes_labels.append(label)
            final_answer = {
                "bboxes": bboxes,
                "bboxes_labels": bboxes_labels,
                "polygons": polygons,
                "polygons_labels": polygons_labels,
            }
        else:
            raise ValueError("Unknown task answer post processing type: {}".format(task_answer_post_processing_type))

        final_answer = {task: final_answer}
        return final_answer


class BoxQuantizer(object):
    def __init__(self, mode, bins):
        self.mode = mode
        self.bins = bins

    def quantize(self, boxes: torch.Tensor, size):
        bins_w, bins_h = self.bins  # Quantization bins.
        size_w, size_h = size  # Original image size.
        size_per_bin_w = size_w / bins_w
        size_per_bin_h = size_h / bins_h
        xmin, ymin, xmax, ymax = boxes.split(1, dim=-1)  # Shape: 4 * [N, 1].

        if self.mode == "floor":
            quantized_xmin = (xmin / size_per_bin_w).floor().clamp(0, bins_w - 1)
            quantized_ymin = (ymin / size_per_bin_h).floor().clamp(0, bins_h - 1)
            quantized_xmax = (xmax / size_per_bin_w).floor().clamp(0, bins_w - 1)
            quantized_ymax = (ymax / size_per_bin_h).floor().clamp(0, bins_h - 1)

        elif self.mode == "round":
            raise NotImplementedError()

        else:
            raise ValueError("Incorrect quantization type.")

        quantized_boxes = torch.cat((quantized_xmin, quantized_ymin, quantized_xmax, quantized_ymax), dim=-1).int()

        return quantized_boxes

    def dequantize(self, boxes: torch.Tensor, size):
        bins_w, bins_h = self.bins  # Quantization bins.
        size_w, size_h = size  # Original image size.
        size_per_bin_w = size_w / bins_w
        size_per_bin_h = size_h / bins_h
        xmin, ymin, xmax, ymax = boxes.split(1, dim=-1)  # Shape: 4 * [N, 1].

        if self.mode == "floor":
            # Add 0.5 to use the center position of the bin as the coordinate.
            dequantized_xmin = (xmin + 0.5) * size_per_bin_w
            dequantized_ymin = (ymin + 0.5) * size_per_bin_h
            dequantized_xmax = (xmax + 0.5) * size_per_bin_w
            dequantized_ymax = (ymax + 0.5) * size_per_bin_h

        elif self.mode == "round":
            raise NotImplementedError()

        else:
            raise ValueError("Incorrect quantization type.")

        dequantized_boxes = torch.cat((dequantized_xmin, dequantized_ymin, dequantized_xmax, dequantized_ymax), dim=-1)

        return dequantized_boxes


class CoordinatesQuantizer(object):
    """
    Quantize coornidates (Nx2)
    """

    def __init__(self, mode, bins):
        self.mode = mode
        self.bins = bins

    def quantize(self, coordinates: torch.Tensor, size):
        bins_w, bins_h = self.bins  # Quantization bins.
        size_w, size_h = size  # Original image size.
        size_per_bin_w = size_w / bins_w
        size_per_bin_h = size_h / bins_h
        assert coordinates.shape[-1] == 2, "coordinates should be shape (N, 2)"
        x, y = coordinates.split(1, dim=-1)  # Shape: 4 * [N, 1].

        if self.mode == "floor":
            quantized_x = (x / size_per_bin_w).floor().clamp(0, bins_w - 1)
            quantized_y = (y / size_per_bin_h).floor().clamp(0, bins_h - 1)

        elif self.mode == "round":
            raise NotImplementedError()

        else:
            raise ValueError("Incorrect quantization type.")

        quantized_coordinates = torch.cat((quantized_x, quantized_y), dim=-1).int()

        return quantized_coordinates

    def dequantize(self, coordinates: torch.Tensor, size):
        bins_w, bins_h = self.bins  # Quantization bins.
        size_w, size_h = size  # Original image size.
        size_per_bin_w = size_w / bins_w
        size_per_bin_h = size_h / bins_h
        assert coordinates.shape[-1] == 2, "coordinates should be shape (N, 2)"
        x, y = coordinates.split(1, dim=-1)  # Shape: 4 * [N, 1].

        if self.mode == "floor":
            # Add 0.5 to use the center position of the bin as the coordinate.
            dequantized_x = (x + 0.5) * size_per_bin_w
            dequantized_y = (y + 0.5) * size_per_bin_h

        elif self.mode == "round":
            raise NotImplementedError()

        else:
            raise ValueError("Incorrect quantization type.")

        dequantized_coordinates = torch.cat((dequantized_x, dequantized_y), dim=-1)

        return dequantized_coordinates


class Florence2PostProcesser(object):
    r"""
    Florence-2 post process for converting text prediction to various tasks results.

    Args:
        config: A dict of configs.
        tokenizer: A tokenizer for decoding text to spans.
        sample config:
            UNIFIED_POST_PROCESS:
                # commom configs
                NUM_BBOX_HEIGHT_BINS: 1000
                NUM_BBOX_WIDTH_BINS: 1000
                COORDINATES_HEIGHT_BINS: 1000
                COORDINATES_WIDTH_BINS: 1000
                # task specific configs, override the common configs
                PRASE_TASKS:
                    - TASK_NAME: 'video_dense_caption'
                      PATTERN: 'r<time_(\d+)><time_(\d+)>([a-zA-Z0-9 ]+)'
                      SCORE_MODE: 'avg_cat_name_scores'
                      NUM_BINS: 100
                    - TASK_NAME: 'od'
                      PATTERN: 'r<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>([a-zA-Z0-9 ]+)'
                      SCORE_MODE: 'avg_cat_name_scores'

    Returns:
        parsed_dict (dict): A dict of parsed results.
    """

    def __init__(self, tokenizer=None):
        parse_tasks = []
        parse_task_configs = {}
        config = self._create_default_config()
        for task in config["PARSE_TASKS"]:
            parse_tasks.append(task["TASK_NAME"])
            parse_task_configs[task["TASK_NAME"]] = task

        self.config = config
        self.parse_tasks = parse_tasks
        self.parse_tasks_configs = parse_task_configs

        self.tokenizer = tokenizer
        if self.tokenizer is not None:
            self.all_special_tokens = set(self.tokenizer.all_special_tokens)

        self.init_quantizers()
        self.black_list_of_phrase_grounding = self._create_black_list_of_phrase_grounding()

    def _create_black_list_of_phrase_grounding(self):
        black_list = {}

        if (
            "phrase_grounding" in self.parse_tasks
            and self.parse_tasks_configs["phrase_grounding"]["FILTER_BY_BLACK_LIST"]
        ):
            black_list = set(
                [
                    "it",
                    "I",
                    "me",
                    "mine",
                    "you",
                    "your",
                    "yours",
                    "he",
                    "him",
                    "his",
                    "she",
                    "her",
                    "hers",
                    "they",
                    "them",
                    "their",
                    "theirs",
                    "one",
                    "oneself",
                    "we",
                    "us",
                    "our",
                    "ours",
                    "you",
                    "your",
                    "yours",
                    "they",
                    "them",
                    "their",
                    "theirs",
                    "mine",
                    "yours",
                    "his",
                    "hers",
                    "its",
                    "ours",
                    "yours",
                    "theirs",
                    "myself",
                    "yourself",
                    "himself",
                    "herself",
                    "itself",
                    "ourselves",
                    "yourselves",
                    "themselves",
                    "this",
                    "that",
                    "these",
                    "those",
                    "who",
                    "whom",
                    "whose",
                    "which",
                    "what",
                    "who",
                    "whom",
                    "whose",
                    "which",
                    "that",
                    "all",
                    "another",
                    "any",
                    "anybody",
                    "anyone",
                    "anything",
                    "each",
                    "everybody",
                    "everyone",
                    "everything",
                    "few",
                    "many",
                    "nobody",
                    "none",
                    "one",
                    "several",
                    "some",
                    "somebody",
                    "someone",
                    "something",
                    "each other",
                    "one another",
                    "myself",
                    "yourself",
                    "himself",
                    "herself",
                    "itself",
                    "ourselves",
                    "yourselves",
                    "themselves",
                    "the image",
                    "image",
                    "images",
                    "the",
                    "a",
                    "an",
                    "a group",
                    "other objects",
                    "lots",
                    "a set",
                ]
            )

        return black_list

    def _create_default_config(self):
        config = {
            "NUM_BBOX_HEIGHT_BINS": 1000,
            "NUM_BBOX_WIDTH_BINS": 1000,
            "BOX_QUANTIZATION_MODE": "floor",
            "COORDINATES_HEIGHT_BINS": 1000,
            "COORDINATES_WIDTH_BINS": 1000,
            "COORDINATES_QUANTIZATION_MODE": "floor",
            "PARSE_TASKS": [
                {"TASK_NAME": "od", "PATTERN": r"([a-zA-Z0-9 ]+)<loc_(\\d+)><loc_(\\d+)><loc_(\\d+)><loc_(\\d+)>"},
                {
                    "TASK_NAME": "ocr",
                    "PATTERN": r"(.+?)<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>",
                    "AREA_THRESHOLD": 0.00,
                },
                {"TASK_NAME": "phrase_grounding", "FILTER_BY_BLACK_LIST": True},
                {
                    "TASK_NAME": "pure_text",
                },
                {
                    "TASK_NAME": "description_with_bboxes",
                },
                {
                    "TASK_NAME": "description_with_polygons",
                },
                {
                    "TASK_NAME": "polygons",
                },
                {
                    "TASK_NAME": "bboxes",
                },
                {
                    "TASK_NAME": "description_with_bboxes_or_polygons",
                },
            ],
        }

        return config

    def init_quantizers(self):
        # we have box_quantizer (od, grounding) and coordinates_quantizer (ocr, referring_segmentation)
        num_bbox_height_bins = self.config.get("NUM_BBOX_HEIGHT_BINS", 1000)
        num_bbox_width_bins = self.config.get("NUM_BBOX_WIDTH_BINS", 1000)
        box_quantization_mode = self.config.get("BOX_QUANTIZATION_MODE", "floor")
        self.box_quantizer = BoxQuantizer(
            box_quantization_mode,
            (num_bbox_width_bins, num_bbox_height_bins),
        )

        num_bbox_height_bins = (
            self.config["COORDINATES_HEIGHT_BINS"]
            if "COORDINATES_HEIGHT_BINS" in self.config
            else self.config.get("NUM_BBOX_HEIGHT_BINS", 1000)
        )
        num_bbox_width_bins = (
            self.config["COORDINATES_WIDTH_BINS"]
            if "COORDINATES_WIDTH_BINS" in self.config
            else self.config.get("NUM_BBOX_WIDTH_BINS", 1000)
        )
        box_quantization_mode = (
            self.config.get("COORDINATES_QUANTIZATION_MODE")
            if "COORDINATES_QUANTIZATION_MODE" in self.config
            else self.config.get("BOX_QUANTIZATION_MODE", "floor")
        )
        self.coordinates_quantizer = CoordinatesQuantizer(
            box_quantization_mode,
            (num_bbox_width_bins, num_bbox_height_bins),
        )

    def decode_with_spans(self, tokenizer, token_ids):
        filtered_tokens = tokenizer.convert_ids_to_tokens(token_ids, skip_special_tokens=False)
        assert len(filtered_tokens) == len(token_ids)

        # To avoid mixing byte-level and unicode for byte-level BPT
        # we need to build string separately for added tokens and byte-level tokens
        # cf. https://github.com/huggingface/transformers/issues/1133
        sub_texts = []
        for token in filtered_tokens:
            if token in self.all_special_tokens:
                sub_texts.append(token)
            else:
                if isinstance(tokenizer, (BartTokenizer, BartTokenizerFast)):
                    sub_text = tokenizer.convert_tokens_to_string([token])
                elif isinstance(tokenizer, (T5Tokenizer, T5TokenizerFast)):
                    # Ref: https://github.com/google/sentencepiece#whitespace-is-treated-as-a-basic-symbol
                    # Note: Do not strip sub_text as it may have functional whitespace
                    sub_text = token.replace("▁", " ")
                else:
                    raise ValueError(f"type {type(tokenizer)} not supported")
                sub_texts.append(sub_text)

        text = ""
        spans = []
        for sub_text in sub_texts:
            span = (len(text), len(text) + len(sub_text))  # [start index, end index).
            text += sub_text
            spans.append(span)

        # Text format:
        # 1. T5Tokenizer/T5TokenizerFast:
        #      "<loc_1><loc_2><loc_3><loc_4> transplanting dog<loc_1><loc_2><loc_3><loc_4> cat</s>"
        #    Equivalent to t5_tokenizer.decode(input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False, spaces_between_special_tokens=False)
        # 2. BartTokenizer (need to double check):
        #      "<s><loc_1><loc_2><loc_3><loc_4>transplanting dog<loc_1><loc_2><loc_3><loc_4>cat</s>"
        #    Equivalent to bart_tokenizer.decode(input_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False, spaces_between_special_tokens=False)
        return text, spans

    def parse_od_from_text_and_spans(self, text, pattern, image_size, phrase_centric=False):
        parsed = list(re.finditer(pattern, text))

        instances = []
        for i in range(len(parsed)):
            # Prepare instance.
            instance = {}

            if phrase_centric:
                bbox_bins = [int(parsed[i].group(j)) for j in range(2, 6)]
            else:
                bbox_bins = [int(parsed[i].group(j)) for j in range(1, 5)]
            instance["bbox"] = self.box_quantizer.dequantize(boxes=torch.tensor(bbox_bins), size=image_size).tolist()

            if phrase_centric:
                instance["cat_name"] = parsed[i].group(1).lower().strip()
            else:
                instance["cat_name"] = parsed[i].group(5).lower().strip()
            instances.append(instance)

        return instances

    def parse_ocr_from_text_and_spans(
        self,
        text,
        pattern,
        image_size,
        area_threshold=-1.0,
    ):
        bboxes = []
        labels = []

        text = text.replace("<s>", "")
        text = text.replace("</s>", "")
        text = text.replace("<eos>", "")
        text = text.replace("<bos>", "")
        text = text.replace("<pad>", "")

        # ocr with regions
        parsed = re.findall(pattern, text)
        instances = []
        image_width, image_height = image_size

        for ocr_line in parsed:
            ocr_content = ocr_line[0]
            quad_box = ocr_line[1:]
            quad_box = [int(i) for i in quad_box]
            quad_box = (
                self.coordinates_quantizer.dequantize(torch.tensor(np.array(quad_box).reshape(-1, 2)), size=image_size)
                .reshape(-1)
                .tolist()
            )

            if area_threshold > 0:
                x_coords = [i for i in quad_box[0::2]]
                y_coords = [i for i in quad_box[1::2]]

                # apply the Shoelace formula
                area = 0.5 * abs(
                    sum(x_coords[i] * y_coords[i + 1] - x_coords[i + 1] * y_coords[i] for i in range(4 - 1))
                )

                if area < (image_width * image_height) * area_threshold:
                    continue

            bboxes.append(quad_box)
            labels.append(ocr_content)
            instances.append(
                {
                    "quad_box": quad_box,
                    "text": ocr_content,
                }
            )
        return instances

    def parse_phrase_grounding_from_text_and_spans(self, text, pattern, image_size):
        # ignore <s> </s> and <pad>
        cur_span = 0
        if text.startswith("<s>"):
            cur_span += 3

        text = text.replace("<s>", "")
        text = text.replace("</s>", "")
        text = text.replace("<eos>", "")
        text = text.replace("<bos>", "")
        text = text.replace("<pad>", "")

        pattern = r"([^<]+(?:<loc_\d+>){4,})"
        phrases = re.findall(pattern, text)

        # pattern should be text pattern and od pattern
        pattern = r"^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_)"
        box_pattern = r"<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>"

        instances = []
        for pharse_text in phrases:
            phrase_text_strip = pharse_text.replace("<ground>", "", 1)
            phrase_text_strip = pharse_text.replace("<obj>", "", 1)

            if phrase_text_strip == "":
                cur_span += len(pharse_text)
                continue

            # Prepare instance.
            instance = {}

            # parse phrase, get string
            phrase = re.search(pattern, phrase_text_strip)
            if phrase is None:
                cur_span += len(pharse_text)
                continue

            # parse bboxes by box_pattern
            bboxes_parsed = list(re.finditer(box_pattern, pharse_text))
            if len(bboxes_parsed) == 0:
                cur_span += len(pharse_text)
                continue

            phrase = phrase.group()
            # remove leading and trailing spaces
            phrase = phrase.strip()

            if phrase in self.black_list_of_phrase_grounding:
                cur_span += len(pharse_text)
                continue

            # a list of list
            bbox_bins = [[int(_bboxes_parsed.group(j)) for j in range(1, 5)] for _bboxes_parsed in bboxes_parsed]
            instance["bbox"] = self.box_quantizer.dequantize(boxes=torch.tensor(bbox_bins), size=image_size).tolist()

            # exclude non-ascii characters
            phrase = phrase.encode("ascii", errors="ignore").decode("ascii")
            instance["cat_name"] = phrase

            instances.append(instance)

        return instances

    def parse_description_with_bboxes_from_text_and_spans(self, text, pattern, image_size, allow_empty_phrase=False):
        # temporary parse solution, split by '.'
        # ignore <s> </s> and <pad>

        text = text.replace("<s>", "")
        text = text.replace("</s>", "")
        text = text.replace("<eos>", "")
        text = text.replace("<bos>", "")
        text = text.replace("<pad>", "")

        if allow_empty_phrase:
            pattern = r"(?:(?:<loc_\d+>){4,})"
        else:
            pattern = r"([^<]+(?:<loc_\d+>){4,})"
        phrases = re.findall(pattern, text)

        # pattern should be text pattern and od pattern
        pattern = r"^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_)"
        box_pattern = r"<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>"

        instances = []
        for pharse_text in phrases:
            phrase_text_strip = pharse_text.replace("<ground>", "", 1)
            phrase_text_strip = pharse_text.replace("<obj>", "", 1)

            if phrase_text_strip == "" and not allow_empty_phrase:
                continue

            # parse phrase, get string
            phrase = re.search(pattern, phrase_text_strip)
            if phrase is None:
                continue

            phrase = phrase.group()
            # remove leading and trailing spaces
            phrase = phrase.strip()

            # parse bboxes by box_pattern
            bboxes_parsed = list(re.finditer(box_pattern, pharse_text))
            if len(bboxes_parsed) == 0:
                continue

            # a list of list
            bbox_bins = [[int(_bboxes_parsed.group(j)) for j in range(1, 5)] for _bboxes_parsed in bboxes_parsed]

            bboxes = self.box_quantizer.dequantize(boxes=torch.tensor(bbox_bins), size=image_size).tolist()

            phrase = phrase.encode("ascii", errors="ignore").decode("ascii")
            for _bboxes in bboxes:
                # Prepare instance.
                instance = {}
                instance["bbox"] = _bboxes
                # exclude non-ascii characters
                instance["cat_name"] = phrase
                instances.append(instance)

        return instances

    def parse_description_with_polygons_from_text_and_spans(
        self,
        text,
        pattern,
        image_size,
        allow_empty_phrase=False,
        polygon_sep_token="<sep>",
        polygon_start_token="<poly>",
        polygon_end_token="</poly>",
        with_box_at_start=False,
    ):
        # ref_seg format: '<expression><x1><y1><x2><y2><><><sep><><><><>'
        # ignore <s> </s> and <pad>

        text = text.replace("<s>", "")
        text = text.replace("</s>", "")
        text = text.replace("<eos>", "")
        text = text.replace("<bos>", "")
        text = text.replace("<pad>", "")

        if allow_empty_phrase:
            pattern = rf"(?:(?:<loc_\d+>|{re.escape(polygon_sep_token)}|{re.escape(polygon_start_token)}|{re.escape(polygon_end_token)}){{4,}})"
        else:
            # [^<]+: This part matches one or more characters that are not the < symbol.
            # The ^ inside the square brackets [] is a negation, meaning it matches anything except <.
            #
            pattern = rf"([^<]+(?:<loc_\d+>|{re.escape(polygon_sep_token)}|{re.escape(polygon_start_token)}|{re.escape(polygon_end_token)}){{4,}})"
        phrases = re.findall(pattern, text)

        phrase_string_pattern = r"^\s*(.*?)(?=<od>|</od>|<box>|</box>|<bbox>|</bbox>|<loc_|<poly>)"
        box_pattern = rf"((?:<loc_\d+>)+)(?:{re.escape(polygon_sep_token)}|$)"

        # one polygons instance is separated by polygon_start_token and polygon_end_token
        polygons_instance_pattern = rf"{re.escape(polygon_start_token)}(.*?){re.escape(polygon_end_token)}"

        instances = []
        for phrase_text in phrases:
            # exclude loc_\d+>
            # need to get span if want to include category score
            phrase_text_strip = re.sub(r"^loc_\d+>", "", phrase_text, count=1)

            # phrase = phrase.replace('<poly>', '')
            # phrase = phrase.replace('poly>', '')

            if phrase_text_strip == "" and not allow_empty_phrase:
                continue

            # parse phrase, get string
            phrase = re.search(phrase_string_pattern, phrase_text_strip)
            if phrase is None:
                continue
            phrase = phrase.group()
            # remove leading and trailing spaces
            phrase = phrase.strip()

            # parse bboxes by box_pattern

            # split by polygon_start_token and polygon_end_token first using polygons_instance_pattern
            if polygon_start_token in phrase_text and polygon_end_token in phrase_text:
                polygons_instances_parsed = list(re.finditer(polygons_instance_pattern, phrase_text))
            else:
                polygons_instances_parsed = [phrase_text]

            for _polygons_instances_parsed in polygons_instances_parsed:
                # Prepare instance.
                instance = {}

                # polygons_parsed= list(re.finditer(box_pattern, phrase_text))
                if isinstance(_polygons_instances_parsed, str):
                    polygons_parsed = list(re.finditer(box_pattern, _polygons_instances_parsed))
                else:
                    polygons_parsed = list(re.finditer(box_pattern, _polygons_instances_parsed.group(1)))
                if len(polygons_parsed) == 0:
                    continue

                # a list of list (polygon)
                bbox = []
                polygons = []
                for _polygon_parsed in polygons_parsed:
                    # group 1: whole <loc_\d+>...</loc_\d+>
                    _polygon = _polygon_parsed.group(1)
                    # parse into list of int
                    _polygon = [int(_loc_parsed.group(1)) for _loc_parsed in re.finditer(r"<loc_(\d+)>", _polygon)]
                    if with_box_at_start and len(bbox) == 0:
                        if len(_polygon) > 4:
                            # no valid bbox prediction
                            bbox = _polygon[:4]
                            _polygon = _polygon[4:]
                        else:
                            bbox = [0, 0, 0, 0]
                    # abandon last element if is not paired
                    if len(_polygon) % 2 == 1:
                        _polygon = _polygon[:-1]

                    # reshape into (n, 2)
                    _polygon = (
                        self.coordinates_quantizer.dequantize(
                            torch.tensor(np.array(_polygon).reshape(-1, 2)), size=image_size
                        )
                        .reshape(-1)
                        .tolist()
                    )
                    # reshape back
                    polygons.append(_polygon)

                instance["cat_name"] = phrase
                instance["polygons"] = polygons
                if len(bbox) != 0:
                    instance["bbox"] = self.box_quantizer.dequantize(
                        boxes=torch.tensor([bbox]), size=image_size
                    ).tolist()[0]

                instances.append(instance)

        return instances

    def __call__(
        self,
        text=None,
        image_size=None,
        parse_tasks=None,
    ):
        """
        Args:
            text: model outputs
            image_size: (width, height)
            parse_tasks: a list of tasks to parse, if None, parse all tasks.

        """
        if parse_tasks is not None:
            if isinstance(parse_tasks, str):
                parse_tasks = [parse_tasks]
            for _parse_task in parse_tasks:
                assert _parse_task in self.parse_tasks, f"parse task {_parse_task} not supported"

        # sequence or text should be provided
        assert text is not None, "text should be provided"

        parsed_dict = {"text": text}

        for task in self.parse_tasks:
            if parse_tasks is not None and task not in parse_tasks:
                continue

            pattern = self.parse_tasks_configs[task].get("PATTERN", None)

            if task == "ocr":
                instances = self.parse_ocr_from_text_and_spans(
                    text,
                    pattern=pattern,
                    image_size=image_size,
                    area_threshold=self.parse_tasks_configs[task].get("AREA_THRESHOLD", 0.0),
                )
                parsed_dict["ocr"] = instances
            elif task == "phrase_grounding":
                instances = self.parse_phrase_grounding_from_text_and_spans(
                    text,
                    pattern=pattern,
                    image_size=image_size,
                )
                parsed_dict["phrase_grounding"] = instances
            elif task == "pure_text":
                parsed_dict["pure_text"] = text
            elif task == "description_with_bboxes":
                instances = self.parse_description_with_bboxes_from_text_and_spans(
                    text,
                    pattern=pattern,
                    image_size=image_size,
                )
                parsed_dict["description_with_bboxes"] = instances
            elif task == "description_with_polygons":
                instances = self.parse_description_with_polygons_from_text_and_spans(
                    text,
                    pattern=pattern,
                    image_size=image_size,
                )
                parsed_dict["description_with_polygons"] = instances
            elif task == "polygons":
                instances = self.parse_description_with_polygons_from_text_and_spans(
                    text,
                    pattern=pattern,
                    image_size=image_size,
                    allow_empty_phrase=True,
                )
                parsed_dict["polygons"] = instances
            elif task == "bboxes":
                instances = self.parse_description_with_bboxes_from_text_and_spans(
                    text,
                    pattern=pattern,
                    image_size=image_size,
                    allow_empty_phrase=True,
                )
                parsed_dict["bboxes"] = instances
            elif task == "description_with_bboxes_or_polygons":
                if "<poly>" in text:
                    # only support either polygons or bboxes, not both at the same time
                    instances = self.parse_description_with_polygons_from_text_and_spans(
                        text,
                        pattern=pattern,
                        image_size=image_size,
                    )
                else:
                    instances = self.parse_description_with_bboxes_from_text_and_spans(
                        text,
                        pattern=pattern,
                        image_size=image_size,
                    )
                parsed_dict["description_with_bboxes_or_polygons"] = instances
            else:
                raise ValueError("task {} is not supported".format(task))

        return parsed_dict