import json import math import re import threading from abc import ABC from textwrap import dedent from typing import Optional, Tuple, Union import numpy as np import torch import torch.nn as nn from PIL import Image, ImageDraw, ImageFont, ImageOps from torchvision import transforms from transformers.image_utils import ImageInput from transformers.processing_utils import ProcessorMixin from transformers.tokenization_utils_base import PreTokenizedInput, TextInput from transformers.utils import logging logger = logging.get_logger(__name__) DS_OCR_PATCH_SIZE = 16 DS_OCR_IMG_TOKEN = "" DS_OCR_DOWNSAMPLE_RATIO = 4 DS_OCR_STOP_STR = "<｜end▁of▁sentence｜>" DS_OCR_TILE_SIZE = 640 DS_OCR_BASE_SIZE = 1024 DEFAULT_OCR_PROMPT = ( "Transcribe the provided document image into a clean Markdown format, preserving the natural reading order." ) DS_OCR_DEFAULT_CHAT_TEMPLATE = dedent( """ {%- for message in messages %} {%- if message['content'] is string %} {{ message['content'].rstrip() }} {%- else %} {%- set ns = namespace(previous_was_image=False) %} {%- for content in message['content'] %} {%- if content['type'] == 'image' %} {{ "" }} {%- set ns.previous_was_image = True %} {%- elif content['type'] == 'text' %} {{- ('\\n' if ns.previous_was_image else '') + content['text'].rstrip() }} {%- set ns.previous_was_image = False %} {%- endif %} {%- endfor %} {%- endif %} {%- if not loop.last %} {{ " " }} {%- endif %} {%- endfor %} """ ).strip() JINA_OCR_CHAT_TEMPLATE = dedent( """ {%- if messages and messages[0]['role'] == 'system' -%} {%- set system_message = messages[0]['content'] -%} {%- set messages = messages[1:] -%} {%- else -%} {%- set system_message = '' -%} {%- endif -%} {{- system_message -}} {%- if system_message %}{{ '\\n' }}{%- endif -%} {%- for message in messages -%} {%- if message['role'] == 'user' -%} {%- set role_prefix = '<|User|>:\n' -%} {%- elif message['role'] == 'assistant' -%} {%- set role_prefix = '<|Assistant|>:\n' -%} {%- else -%} {%- set role_prefix = '<|' ~ message['role']|capitalize ~ '|>:\n' -%} {%- endif -%} {%- if message['content'] is string -%} {{- role_prefix + message['content'].rstrip() -}} {%- else -%} {%- set ns = namespace(prefix_printed=False, last_was_image=False) -%} {%- for content in message['content'] -%} {%- set is_image = content.get('type') in ['image', 'image_pil', 'image_url'] -%} {%- set content_text = '' if is_image else content.get('text', '').rstrip() -%} {%- if ns.last_was_image and not is_image %}{{ '\\n' }}{%- endif -%} {%- if not ns.prefix_printed and (message['role'] == 'user' or not is_image) -%} {{- role_prefix -}} {%- set ns.prefix_printed = True -%} {%- endif -%} {{- content_text -}} {%- set ns.last_was_image = is_image -%} {%- endfor -%} {%- endif -%} {%- if not loop.last %}{{ '\\n' }}{%- endif -%} {%- endfor -%} {%- if messages and messages|length > 0 and messages[messages|length - 1]['role'] == 'assistant' -%} {{- eos_token|default('') -}} {%- elif add_generation_prompt -%} {%- if messages %}{{ '\\n' }}{%- endif -%} {{- '<|Assistant|>:\n' -}} {%- endif -%} """ ).strip() # Known element types - if label is not one of these, it's an image with alt text KNOWN_ELEMENT_TYPES = { "title", "text", "paragraph", "button", "link", "icon", "header", "footer", "table", "list", "code", "formula", "caption", "label", "input", "checkbox", "radio", "dropdown", "menu", "navigation", "sidebar", "logo", "banner", "card", } # Regex pattern for matching ref/det tags REF_DET_PATTERN = re.compile( r"(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)", re.DOTALL, ) # Regex pattern for matching markdown image syntax: ![alt_text](**/x1_y1_x2_y2.png|jpg|jpeg) MARKDOWN_IMAGE_PATTERN = re.compile( r"!\[([^\]]*)\]\([^)]*?(\d+)_(\d+)_(\d+)_(\d+)\.(?:png|jpg|jpeg)\)", re.DOTALL, ) def normalize_markdown_images_to_ref_det(text: str) -> str: """ Normalize markdown image syntax to ref/det format. Converts: ![alt_description](**/x1_y1_x2_y2.png) To: <|ref|>alt_description<|/ref|><|det|>[[x1,y1,x2,y2]]<|/det|> If alt_description is empty, uses "image" as the default label. Args: text: Input text containing markdown image syntax Returns: Text with markdown images converted to ref/det format """ def replace_match(match: re.Match) -> str: alt_text = match.group(1).strip() x1 = match.group(2) y1 = match.group(3) x2 = match.group(4) y2 = match.group(5) # Use "image" as default label if alt text is empty label = alt_text if alt_text else "image" return f"<|ref|>{label}<|/ref|><|det|>[[{x1},{y1},{x2},{y2}]]<|/det|>" return MARKDOWN_IMAGE_PATTERN.sub(replace_match, text) def parse_coords(coords_str: str) -> list[list[int]]: """ Parse a coordinate string into a list of bounding boxes. Args: coords_str: JSON string of coordinates, e.g. "[[100,200,300,400]]" or "[[100,200,300,400],[500,600,700,800]]" Returns: List of [x1, y1, x2, y2] integer coordinates. Empty list if parsing fails. Note: All coordinates are normalized into 1000 bins (range [0, 999]). """ try: coords_list = json.loads(coords_str) return [[int(a), int(b), int(c), int(d)] for a, b, c, d in coords_list] except Exception: return [] def normalize_to_pixel(x: int, y: int, image_width: int, image_height: int) -> tuple[int, int]: """ Convert normalized coordinates (0-999) to actual pixel coordinates. Args: x: Normalized x coordinate (0-999) y: Normalized y coordinate (0-999) image_width: Actual image width in pixels image_height: Actual image height in pixels Returns: Tuple of (pixel_x, pixel_y) Note: All coordinates are normalized into 1000 bins (range [0, 999]). Formula: pixel = int(normalized / 999 * dimension) """ return int(x / 999 * image_width), int(y / 999 * image_height) def coords_to_filename(x1: int, y1: int, x2: int, y2: int, ext: str = "png") -> str: """ Generate a filename from bounding box coordinates. Args: x1, y1: Top-left corner (normalized 0-999) x2, y2: Bottom-right corner (normalized 0-999) ext: File extension (default: "png") Returns: Filename string like "100_200_300_400.png" """ return f"{x1}_{y1}_{x2}_{y2}.{ext}" def is_image_label(label: str) -> bool: """ Determine if a label represents an image (i.e., not a known element type). Args: label: The label string from a ref tag Returns: True if this is an image (label is NOT a known element type), False if it's a known UI element type. Note: When True, the label serves as the alt text for the image in markdown. """ return label.lower() not in KNOWN_ELEMENT_TYPES def parse_refs(text: str) -> list[dict]: """ Extract labeled bounding box references from text. Matches patterns like: <|ref|>label<|/ref|><|det|>[[x1,y1,x2,y2]]<|/det|> # single bounding box <|ref|>label<|/ref|><|det|>[[x1,y1,x2,y2],[x1,y1,x2,y2]]<|/det|> # multiple bounding boxes Coordinate System: All coordinates are normalized into 1000 bins (range [0, 999]), representing positions relative to the original image dimensions. - (x1, y1): Top-left corner of the bounding box - (x2, y2): Bottom-right corner of the bounding box To convert to actual pixel coordinates, use normalize_to_pixel(). Visual representation: (0,0) ─────────────────────────────► x (width) │ │ (x1,y1) ┌────────────┐ │ │ │ │ │ element │ │ │ │ │ └────────────┘ (x2,y2) │ ▼ y (height) Image detection: If the label is NOT a known element type (see KNOWN_ELEMENT_TYPES), it is treated as an image where the label serves as the alt text. Args: text: The text to parse for refs Returns: List of dicts with keys: - "label": str - The identifier between <|ref|> and <|/ref|> - "coords_literal": str - The raw coordinate string between <|det|> and <|/det|> - "triplet": [full_match, label, coords_literal] - Structured representation - "coords": list[list[int]] - Parsed coordinates as [[x1,y1,x2,y2], ...] - "span": (start, end) - Character positions in original text - "is_image": bool - True if label is NOT a known element type """ refs = [] for m in REF_DET_PATTERN.finditer(text): full_match = m.group(1) label = m.group(2) coords_literal = m.group(3) refs.append( { "label": label, "coords_literal": coords_literal, "triplet": [full_match, label, coords_literal], "coords": parse_coords(coords_literal), "span": (m.start(), m.end()), "is_image": is_image_label(label), } ) return refs def re_match(text): """ Extract labeled bounding box references from text, split by image vs non-image. This is a convenience wrapper around parse_refs() that returns results grouped by whether they are images or known element types. See parse_refs() for full documentation on coordinate system and pattern matching. Returns: tuple: (matches, matches_image, matches_other) - matches: list of tuples (full_match, label, coords_literal) - matches_image: list of tuples (full_match, label, coords_literal) for images - matches_other: list of full_match strings for non-image elements """ refs = parse_refs(text) matches = [tuple(r["triplet"]) for r in refs] matches_image = [tuple(r["triplet"]) for r in refs if r["is_image"]] matches_other = [r["triplet"][0] for r in refs if not r["is_image"]] return matches, matches_image, matches_other def extract_coordinates_and_label(ref_text): """ Extract label and coordinate list from a ref_text triplet. Args: ref_text: tuple/list of (full_match, label, coords_literal) - full_match: the complete matched string - label: the identifier (or alt text for images) - coords_literal: the coordinate string, e.g. "[[100,200,300,400]]" Returns: tuple: (label_type, coords_list) where coords_list is a list of [x1, y1, x2, y2] boxes None: if parsing fails """ try: label_type = ref_text[1] coords_list = parse_coords(ref_text[2]) if not coords_list: return None return (label_type, coords_list) except Exception as e: print(e) return None def text_encode(tokenizer, text: str, bos: bool = True, eos: bool = False): """ Encode text using the tokenizer with optional BOS/EOS tokens. Args: tokenizer: The tokenizer to use text: Text to encode bos: Whether to prepend BOS token (default: True) eos: Whether to append EOS token (default: False) Returns: List of token IDs """ t = tokenizer.encode(text, add_special_tokens=False) bos_id = 0 eos_id = 1 if bos: t = [bos_id] + t if eos: t = t + [eos_id] return t # ============================================================================= # Image Transform Utilities # ============================================================================= def normalize_transform(mean, std): """ Create a normalization transform for image tensors. Args: mean: Mean values for normalization (per channel) std: Standard deviation values for normalization (per channel) Returns: transforms.Normalize instance or None if both mean and std are None """ if mean is None and std is None: transform = None elif mean is None and std is not None: mean = [0.0] * len(std) transform = transforms.Normalize(mean=mean, std=std) elif mean is not None and std is None: std = [1.0] * len(mean) transform = transforms.Normalize(mean=mean, std=std) else: transform = transforms.Normalize(mean=mean, std=std) return transform def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size): """ Find the closest aspect ratio from a set of target ratios. Args: aspect_ratio: The aspect ratio to match target_ratios: Set of (width_ratio, height_ratio) tuples to choose from width: Image width height: Image height image_size: Base tile size Returns: Tuple (width_ratio, height_ratio) representing the best match """ best_ratio_diff = float("inf") best_ratio = (1, 1) area = width * height for ratio in target_ratios: target_aspect_ratio = ratio[0] / ratio[1] ratio_diff = abs(aspect_ratio - target_aspect_ratio) if ratio_diff < best_ratio_diff: best_ratio_diff = ratio_diff best_ratio = ratio elif ratio_diff == best_ratio_diff: if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]: best_ratio = ratio return best_ratio def dynamic_preprocess(image, min_num=2, max_num=9, image_size=640, use_thumbnail=False): """ Dynamically preprocess an image by tiling it based on aspect ratio. Splits the image into multiple tiles based on the closest matching aspect ratio from a set of valid tile arrangements. Args: image: PIL Image to preprocess min_num: Minimum number of tiles (default: 2) max_num: Maximum number of tiles (default: 9) image_size: Size of each tile (default: 640) use_thumbnail: Whether to append a thumbnail of the full image (default: False) Returns: Tuple of (processed_images, target_aspect_ratio) where: - processed_images: List of PIL Image tiles - target_aspect_ratio: Tuple (width_tiles, height_tiles) """ orig_width, orig_height = image.size aspect_ratio = orig_width / orig_height # calculate the existing image aspect ratio target_ratios = set( (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if i * j <= max_num and i * j >= min_num ) target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1]) # find the closest aspect ratio to the target target_aspect_ratio = find_closest_aspect_ratio(aspect_ratio, target_ratios, orig_width, orig_height, image_size) # calculate the target width and height target_width = image_size * target_aspect_ratio[0] target_height = image_size * target_aspect_ratio[1] blocks = target_aspect_ratio[0] * target_aspect_ratio[1] # resize the image resized_img = image.resize((target_width, target_height)) processed_images = [] for i in range(blocks): box = ( (i % (target_width // image_size)) * image_size, (i // (target_width // image_size)) * image_size, ((i % (target_width // image_size)) + 1) * image_size, ((i // (target_width // image_size)) + 1) * image_size, ) # split the image split_img = resized_img.crop(box) processed_images.append(split_img) assert len(processed_images) == blocks if use_thumbnail and len(processed_images) != 1: thumbnail_img = image.resize((image_size, image_size)) processed_images.append(thumbnail_img) return processed_images, target_aspect_ratio class BaseTransform(ABC): """Abstract base class for image transforms.""" def set_rng(self, *args, **kwargs): pass def __call__(self, *args, **kwargs) -> torch.Tensor: pass @property def default_shape(self): raise NotImplementedError class BasicImageTransform(BaseTransform): """ Basic image transform that converts PIL images to normalized tensors. Args: mean: Mean values for normalization (default: (0.5, 0.5, 0.5)) std: Standard deviation values for normalization (default: (0.5, 0.5, 0.5)) normalize: Whether to apply normalization (default: True) """ def __init__( self, mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), normalize: bool = True, ): self.mean = mean self.std = std transform_pipelines = [transforms.ToTensor()] normalize_op = normalize_transform(mean, std) if normalize else nn.Identity() if normalize_op is not None: transform_pipelines.append(normalize_op) self.transform = transforms.Compose(transform_pipelines) def __call__(self, x): x = self.transform(x) return x class DeepseekOCRProcessor(ProcessorMixin): attributes = ["tokenizer"] # image_processor_class = "DeepseekOcrImageProcessorFast" tokenizer_class = "AutoTokenizer" def __init__( self, tokenizer, use_fast: Optional[bool] = None, crop_mode: bool = True, base_size: int = DS_OCR_BASE_SIZE, image_size: int = DS_OCR_TILE_SIZE, **kwargs, ): """ Initialize the DeepseekOCRProcessor. Args: tokenizer: A HuggingFace tokenizer that must contain the special image token (defaults to ``DS_OCR_IMG_TOKEN = ""``). use_fast: Whether the tokenizer is a fast tokenizer. If provided, it must match the actual ``tokenizer.is_fast`` value. Kept for API compatibility. crop_mode: Default resolution strategy for image preprocessing. - ``True`` (Gundam / dynamic resolution): processes a global view at ``base_size`` **and** up to 9 local tiles at ``image_size``. Best for high-resolution documents where fine spatial detail matters. - ``False`` (native / fixed resolution): resizes the image to a single ``image_size × image_size`` tensor. Simpler and faster; used by smaller model variants (Tiny / Small / Base / Large). Can be overridden per-call via the ``crop_mode`` argument of ``__call__``. **kwargs: Additional keyword arguments forwarded to ``ProcessorMixin``. Supports ``chat_template`` (``"jina"``, ``"deepseek"``, or a full template string). """ if hasattr(tokenizer, "image_token"): self.image_token = tokenizer.image_token else: self.image_token = kwargs.pop("image_token", DS_OCR_IMG_TOKEN) self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token) if self.image_token_id is None: raise ValueError( f"The tokenizer does not contain the special image token `{self.image_token}`. " "Please make sure it is added to the vocabulary before instantiating the processor." ) # WARNING: use_fast is parameter for tokenizer, but here we keep it for compatibility. self.use_fast = use_fast if self.use_fast is not None: assert self.use_fast == tokenizer.is_fast, ( f"The tokenizer fast option use_fast={self.use_fast} does not match the actual " f"tokenizer type (is_fast={tokenizer.is_fast})." ) if "chat_template" not in kwargs: kwargs["chat_template"] = JINA_OCR_CHAT_TEMPLATE if kwargs["chat_template"] == "jina": kwargs["chat_template"] = JINA_OCR_CHAT_TEMPLATE elif kwargs["chat_template"] == "deepseek": kwargs["chat_template"] = DS_OCR_DEFAULT_CHAT_TEMPLATE # else: # raise ValueError(f"Invalid chat template: {kwargs['chat_template']}") tokenizer.chat_template = kwargs["chat_template"] super().__init__(tokenizer, **kwargs) self.image_transform = BasicImageTransform(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5), normalize=True) self.crop_mode = crop_mode self.base_size = base_size self.image_size = image_size def __call__( self, images: Optional[ImageInput] = None, text: Optional[Union[TextInput, PreTokenizedInput, list[TextInput], list[PreTokenizedInput]]] = None, base_size: Optional[int] = None, image_size: Optional[int] = None, crop_mode: Optional[bool] = None, return_tensors: str = "pt", thread_lock: Optional[threading.Lock] = None, **kwargs, ): """ Preprocess a single (image, text) pair into model-ready tensors. Args: images: A list containing exactly one PIL ``Image``. Only single-image inputs are supported. text: The prompt string (or a single-element list thereof). Any markdown image syntax (``![alt](path/x1_y1_x2_y2.png)``) is automatically converted to the ``<|ref|>…<|/ref|><|det|>…<|/det|>`` format before tokenization. base_size: Side length in pixels of the **global view** tile, used only when ``crop_mode=True``. The entire input image is padded/resized to ``base_size × base_size`` and encoded as a single high-level overview. The number of image tokens produced is ``(ceil(base_size / PATCH_SIZE / DOWNSAMPLE_RATIO))² + 1`` per row/col. Default: ``DS_OCR_BASE_SIZE = 1024``. image_size: Side length in pixels of each **local tile**, used by ``dynamic_preprocess``. - When ``crop_mode=True``: the image is split into 2–9 tiles of ``image_size × image_size`` arranged in the grid layout that best preserves the original aspect ratio (e.g. 2×3, 1×4, …). Each tile generates its own token sequence appended after the global-view tokens. - When ``crop_mode=False``: the image is simply resized/padded to ``image_size × image_size`` and this is the only resolution used. Default: ``DS_OCR_TILE_SIZE = 640``. crop_mode: Resolution strategy for this call. Overrides the instance-level ``self.crop_mode`` when provided. - ``True`` — dual-stream: global view (``base_size``) + local tiles (``image_size``). Captures both layout context and fine detail. - ``False`` — single-stream: image resized to ``image_size`` only. Lower memory footprint; suitable for smaller model variants. Defaults to ``None``, which falls back to ``self.crop_mode``. return_tensors: Output tensor format. Only ``"pt"`` (PyTorch) is supported. thread_lock: An optional ``threading.Lock`` to serialise tokenizer calls. The HuggingFace tokenizer is not thread-safe; pass a shared lock when calling this processor from multiple threads concurrently. **kwargs: Unused; accepted for forward-compatibility. Returns: A dict with the following keys: - ``"input_ids"`` (``LongTensor``, shape ``[1, seq_len]``): Token IDs for the full prompt, with image placeholder tokens interleaved. - ``"images"`` (``tuple[Tensor, Tensor]``): ``(images_crop, images_ori)`` where ``images_ori`` holds the global view(s) (shape ``[n_views, 3, base_size, base_size]``) and ``images_crop`` holds the local tiles (shape ``[n_tiles, 3, image_size, image_size]``), or an empty tensor when ``crop_mode=False``. - ``"images_seq_mask"`` (``BoolTensor``, shape ``[seq_len]``): ``True`` at every position occupied by an image placeholder token. - ``"images_spatial_crop"`` (``LongTensor``, shape ``[n_views, 2]``): Each row is ``[width_tile_count, height_tile_count]`` for the corresponding image, encoding the grid layout chosen by ``dynamic_preprocess``. """ assert return_tensors == "pt", "Only PyTorch tensors are supported for DeepseekOCRProcessor." if isinstance(text, list): assert len(text) == 1, "Only single text input is supported." text = text[0] assert len(images) == 1, "Only single image input is supported." base_size = self.base_size if base_size is None else base_size image_size = self.image_size if image_size is None else image_size # Normalize markdown image syntax to ref/det format assert isinstance(text, str), "Text must be a string at this point." text = normalize_markdown_images_to_ref_det(text) input_ids, images_seq_mask, images_crop, images_ori, images_spatial_crop = self.preprocess_prompt_and_image( prompt=text, image=images[0], base_size=base_size, image_size=image_size, crop_mode=self.crop_mode if crop_mode is None else crop_mode, thread_lock=thread_lock, ) batched_ids = input_ids.unsqueeze(0) return { "input_ids": batched_ids, "attention_mask": torch.ones_like(batched_ids), "images": (images_crop, images_ori), "images_seq_mask": images_seq_mask, "images_spatial_crop": images_spatial_crop, } @staticmethod def to_device(batch: dict, device): """Move processor outputs (including the ``images`` tensor tuple) to ``device``.""" def _move(value): if torch.is_tensor(value): return value.to(device) if isinstance(value, (list, tuple)): return type(value)(_move(v) for v in value) return value return {key: _move(value) for key, value in batch.items()} def prepare_ocr_inputs( self, image: Image.Image, prompt: Optional[str] = None, device=None, **kwargs, ) -> dict: """Apply the chat template and return generation-ready tensors.""" prompt = DEFAULT_OCR_PROMPT if prompt is None else prompt conversation = [ { "role": "user", "content": [ {"type": "image", "image": image}, {"type": "text", "text": prompt}, ], } ] text = self.apply_chat_template(conversation, add_generation_prompt=True) inputs = self(text=[text], images=[image], return_tensors="pt", **kwargs) # generate() expects a batch of (crop, ori) pairs for zip(*images). if isinstance(inputs.get("images"), tuple) and len(inputs["images"]) == 2: inputs["images"] = [inputs["images"]] if device is not None: inputs = self.to_device(inputs, device) return inputs def decode_ocr(self, sequences, input_ids: torch.Tensor) -> str: """Decode generated tokens after the prompt. DeepSeek-OCR training decodes with ``skip_special_tokens=False`` because ````, ``<|ref|>``, and EOS are all marked special — skipping them turns an EOS/image-only continuation into an empty string. """ if hasattr(sequences, "sequences"): sequences = sequences.sequences generated = sequences[0][input_ids.shape[-1] :] text = self.tokenizer.decode(generated, skip_special_tokens=False) for tok in ( getattr(self.tokenizer, "eos_token", None), getattr(self.tokenizer, "pad_token", None), getattr(self.tokenizer, "bos_token", None), self.image_token, ): if tok: text = text.replace(tok, "") text = text.strip() if not text and generated.numel() > 0: unique = generated.detach().cpu().unique().tolist() logger.warning( "decode_ocr is empty after stripping specials from %s new tokens " "(unique ids=%s). generate() likely never produced text — check " "eos_token_id and that the model is not repeating /EOS.", generated.numel(), unique[:16], ) return text def preprocess_prompt_and_image( self, prompt: str, image: Image.Image, base_size: int, image_size: int, crop_mode: bool, thread_lock: Optional[threading.Lock] = None, ): tokenized_image, images_list, images_crop_list, images_spatial_crop = self.preprocess_image( image=image, base_size=base_size, image_size=image_size, crop_mode=crop_mode, ) prompt_chunks = prompt.split(self.image_token) # Protect tokenizer calls with thread lock if provided (tokenizer is not thread-safe) if thread_lock is not None: with thread_lock: pre_tokens, post_tokens = [ text_encode(self.tokenizer, pc, bos=False, eos=False) for pc in prompt_chunks ] else: pre_tokens, post_tokens = [text_encode(self.tokenizer, pc, bos=False, eos=False) for pc in prompt_chunks] tokenized_str = [*pre_tokens, *tokenized_image, *post_tokens] input_ids = torch.LongTensor(tokenized_str) images_crop = torch.zeros(0, 3, image_size, image_size) if len(images_list) == 0: ori_size = base_size if crop_mode else image_size images_ori = torch.zeros((1, 3, ori_size, ori_size)) images_spatial_crop = torch.zeros((1, 2), dtype=torch.long) else: images_ori = torch.stack(images_list, dim=0) images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) if images_crop_list: images_crop = torch.stack(images_crop_list, dim=0) images_seq_mask = (input_ids == self.image_token_id).to(torch.bool) return input_ids, images_seq_mask, images_crop, images_ori, images_spatial_crop def tokenize_prompt_with_image_count( self, prompt: str, image_token_count: int, thread_lock: Optional[threading.Lock] = None, ) -> torch.LongTensor: """Tokenize text while reusing an already-processed image token count.""" prompt = normalize_markdown_images_to_ref_det(prompt) prompt_chunks = prompt.split(self.image_token) if len(prompt_chunks) != 2: raise ValueError(f"Expected exactly one {self.image_token} placeholder, found {len(prompt_chunks) - 1}.") if thread_lock is not None: with thread_lock: pre_tokens, post_tokens = [ text_encode(self.tokenizer, chunk, bos=False, eos=False) for chunk in prompt_chunks ] else: pre_tokens, post_tokens = [ text_encode(self.tokenizer, chunk, bos=False, eos=False) for chunk in prompt_chunks ] return torch.LongTensor([*pre_tokens, *([self.image_token_id] * image_token_count), *post_tokens]) def preprocess_image( self, image: Image.Image, base_size: int, image_size: int, crop_mode: bool, ): """ Convert a single PIL image into token IDs and normalized tensors. Depending on ``crop_mode`` two pipelines are available: **crop_mode = True** (Gundam / dynamic resolution) 1. The full image is padded to ``base_size × base_size`` → *global view*. 2. If the image is larger than ``image_size`` in either dimension it is also passed through ``dynamic_preprocess``, which selects the grid layout (``width_crop_num × height_crop_num``, 2–9 tiles) whose aspect ratio is closest to the original, resizes the image to fit that grid, and slices it into ``image_size × image_size`` tiles → *local tiles*. 3. Image tokens: global-view tokens (derived from ``base_size``) followed by patch tokens for each local tile (derived from ``image_size``). **crop_mode = False** (native / fixed resolution) The image is resized directly to ``image_size × image_size``. No tiling is performed and only global-view tokens are emitted. Args: image: Input PIL image in any mode/size. base_size: Target size for the global view (used in ``crop_mode=True``). See ``__call__`` for details. image_size: Target tile size for local crops (``crop_mode=True``) or the single target resolution (``crop_mode=False``). See ``__call__`` for details. crop_mode: Whether to apply dynamic tiling. See ``__call__`` for details. Returns: A 4-tuple ``(tokenized_image, images_list, images_crop_list, images_spatial_crop)``: - ``tokenized_image`` (``list[int]``): Image placeholder token IDs to be spliced into the prompt token sequence. - ``images_list`` (``list[Tensor]``): Normalized global-view tensor(s), each of shape ``[3, base_size, base_size]``. - ``images_crop_list`` (``list[Tensor]``): Normalized local-tile tensor(s), each of shape ``[3, image_size, image_size]``. Empty when no tiling occurs. - ``images_spatial_crop`` (``list[[int, int]]``): ``[[width_crop_num, height_crop_num]]`` recording the tile grid dimensions. """ images_list, images_crop_list = [], [] images_spatial_crop = [] if image.mode != "RGB": image = image.convert("RGB") width_crop_num, height_crop_num = 1, 1 if crop_mode: # Dynamic resolution (Gundam) if not (image.size[0] <= image_size and image.size[1] <= image_size): images_crop_raw, (width_crop_num, height_crop_num) = dynamic_preprocess(image, image_size=image_size) # Process the global view global_view = self.pad_img(image, base_size) images_list.append(self.image_transform(global_view)) if width_crop_num > 1 or height_crop_num > 1: # Process the local views for i in range(len(images_crop_raw)): images_crop_list.append(self.image_transform(images_crop_raw[i])) num_queries = self.compute_n_queries(image_size) num_queries_base = self.compute_n_queries(base_size) # Add image tokens tokenized_image = self.image_tokens(num_queries_base) if width_crop_num > 1 or height_crop_num > 1: tokenized_image += self.patch_image_tokens(num_queries, width_crop_num, height_crop_num) else: # Native resolution (Tiny/Small/Base/Large) # Process the global view if image_size <= 640: # Tiny/Small mode: image fits within target, directly resize image = image.resize((image_size, image_size)) global_view = self.pad_img(image, image_size) images_list.append(self.image_transform(global_view)) # Add image tokens""" num_queries = self.compute_n_queries(image_size) tokenized_image = self.image_tokens(num_queries) images_spatial_crop.append([width_crop_num, height_crop_num]) return tokenized_image, images_list, images_crop_list, images_spatial_crop def compute_n_queries( self, image_size: int, ): """ Compute the number of image query tokens per row (or column) for a square tile. The vision encoder first divides the image into patches of size ``DS_OCR_PATCH_SIZE`` (16 px), then a pixel-shuffle / downsample operation reduces the spatial grid by ``DS_OCR_DOWNSAMPLE_RATIO`` (4×) in each dimension. The result is the 1-D query count used to build the 2-D grid of image placeholder tokens. Formula:: n = ceil((image_size // PATCH_SIZE) / DOWNSAMPLE_RATIO) = ceil((image_size // 16) / 4) Examples: - ``image_size=640`` → ``ceil(40 / 4)`` = **10** queries per side - ``image_size=1024`` → ``ceil(64 / 4)`` = **16** queries per side Args: image_size: Side length of the (square) tile in pixels. Returns: Number of query tokens along one spatial dimension. """ num_queries = math.ceil((image_size // DS_OCR_PATCH_SIZE) / DS_OCR_DOWNSAMPLE_RATIO) return num_queries def compute_valid_img_tokens( self, base_size: int, ratio: float, ): if base_size == 1024: return int(256 * ratio) elif base_size == 1280: return int(400 * ratio) elif base_size == 640: return int(100 * 1) elif base_size == 512: return int(64 * 1) raise ValueError("Unsupported base_size for valid image tokens computation.") def pad_img(self, img, size): """ Pad (and/or resize) a PIL image to a square canvas of ``size × size`` pixels. Uses ``ImageOps.pad``, which scales the image down to fit within the target square while preserving aspect ratio, then fills the remaining border with the mean colour of the normalisation transform (converted from [0, 1] to [0, 255]). This avoids introducing out-of-distribution black or white borders. Args: img: Input PIL image. size: Target side length in pixels for the output square canvas. Returns: A new PIL image of size ``(size, size)``. """ color = tuple(int(x * 255) for x in self.image_transform.mean) return ImageOps.pad( img, (size, size), color=color, ) def image_tokens(self, num_queries) -> list[int]: """ Build image placeholder token IDs for one square view. The model encodes this as ``num_queries`` rows of visual features, appends one newline token to each row, then appends one final view separator token. Args: num_queries: Number of query tokens per spatial dimension (from ``compute_n_queries``). Returns: A flat ``list[int]`` of length ``num_queries * (num_queries + 1) + 1``. """ base = [self.image_token_id] * num_queries + [self.image_token_id] return base * num_queries + [self.image_token_id] def patch_image_tokens(self, num_queries, width_crop_num, height_crop_num): """ Build the flat list of image placeholder token IDs for all local tiles. Generates tokens for the full ``width_crop_num × height_crop_num`` tile grid produced by ``dynamic_preprocess``. Each tile contributes ``num_queries`` tokens along its dimension, plus one separator token per row/column. Layout (width=W tiles, height=H tiles, q queries per tile side):: row length = num_queries * width_crop_num + 1 (one separator per row) n_rows = num_queries * height_crop_num (no trailing separator) Args: num_queries: Number of query tokens per spatial dimension per tile (from ``compute_n_queries`` called with ``image_size``). width_crop_num: Number of tile columns in the grid (``target_aspect_ratio[0]``). height_crop_num: Number of tile rows in the grid (``target_aspect_ratio[1]``). Returns: A flat ``list[int]`` of image token IDs covering the entire local-tile grid. """ base = [self.image_token_id] * (num_queries * width_crop_num) + [self.image_token_id] return base * (num_queries * height_crop_num) def build_ref_texts_from_output(self, text: str): """Extract triplets from parsed refs for backward compatibility.""" return [r["triplet"] for r in parse_refs(text)] def visualize_results(self, image, ref_texts): if isinstance(image, torch.Tensor): image = Image.fromarray((image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)) width, height = image.size img_draw = image.copy() draw = ImageDraw.Draw(img_draw) overlay = Image.new("RGBA", img_draw.size, (0, 0, 0, 0)) draw2 = ImageDraw.Draw(overlay) try: font = ImageFont.load_default() except Exception: font = None for ref in ref_texts: label_type, boxes = extract_coordinates_and_label(ref) color = ( np.random.randint(0, 200), np.random.randint(0, 200), np.random.randint(0, 255), ) color_a = color + (20,) for orig_x1, orig_y1, orig_x2, orig_y2 in boxes: x1, y1 = normalize_to_pixel(orig_x1, orig_y1, width, height) x2, y2 = normalize_to_pixel(orig_x2, orig_y2, width, height) if label_type == "title": draw.rectangle([x1, y1, x2, y2], outline=color, width=4) draw2.rectangle([x1, y1, x2, y2], fill=color_a) else: draw.rectangle([x1, y1, x2, y2], outline=color, width=2) draw2.rectangle([x1, y1, x2, y2], fill=color_a) if font: text_x, text_y = x1, max(0, y1 - 15) textbox_width, textbox_height = draw.textbbox((0, 0), label_type, font=font)[2:] draw.rectangle( [ text_x, text_y, text_x + textbox_width, text_y + textbox_height, ], fill=(255, 255, 255, 30), ) draw.text((text_x, text_y), label_type, font=font, fill=color) img_draw.paste(overlay, (0, 0), overlay) return img_draw def extract_markdown_and_crops( self, text: str, image, image_placeholder_fmt="![{alt_text}](images/{x1}_{y1}_{x2}_{y2}.png)", cleanup=True, ): """ Extract markdown text and image crops from model output. Parses ref/det tags in the output text and: - Replaces image references with markdown image syntax (e.g., `![alt](images/x1_y1_x2_y2.png)`) - Removes non-image element references (e.g., title, text, button) from the output - Crops the corresponding image regions for each detected image Image detection: A ref is considered an image if its label is NOT one of the known UI element types (title, text, button, etc.). The label then serves as the alt text. Example: Input: "<|ref|>A chart showing sales<|/ref|><|det|>[[100,200,300,400]]<|/det|>" Output: "![A chart showing sales](images/100_200_300_400.png)" Args: text: Model output text containing `<|ref|>...<|/ref|><|det|>...<|/det|>` tags image: Source image as PIL Image or torch.Tensor (C, H, W) image_placeholder_fmt: Format string for markdown image syntax. Available placeholders: {alt_text}, {x1}, {y1}, {x2}, {y2} cleanup: If True, replace common LaTeX symbols (e.g., \\coloneqq -> :=) Returns: tuple: (markdown_text, crops) - markdown_text: Processed text with image refs converted to markdown - crops: List of (cropped_image, alt_text, filename) tuples for each image """ refs = parse_refs(text) pieces = [] last = 0 crops = [] if isinstance(image, torch.Tensor): image = Image.fromarray((image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)) W, H = image.size for r in refs: start, end = r["span"] pieces.append(text[last:start]) if r["is_image"]: alt_text = r["label"] for x1, y1, x2, y2 in r["coords"]: X1, Y1 = normalize_to_pixel(x1, y1, W, H) X2, Y2 = normalize_to_pixel(x2, y2, W, H) if X2 > X1 and Y2 > Y1: filename = coords_to_filename(x1, y1, x2, y2) crops.append((image.crop((X1, Y1, X2, Y2)), alt_text, filename)) pieces.append(image_placeholder_fmt.format(alt_text=alt_text, x1=x1, y1=y1, x2=x2, y2=y2)) # non-image refs are dropped from text last = end pieces.append(text[last:]) md = "".join(pieces) if cleanup: md = md.replace("\\coloneqq", ":=").replace("\\eqqcolon", "=:") return md, crops