from __future__ import annotations import threading from dataclasses import dataclass, field from typing import Any from PIL import Image from ocr_studio.config import ( LANGUAGE_BOTH, LANGUAGE_EN, LANGUAGE_FA, MAX_NEW_TOKENS, MODE_COMPARE, MODE_DOCUMENT, MODE_FIELDS, MODE_MARKDOWN, MODE_PRECISE, MODE_TABLE, OCR_MAX_PIXELS, SPOTTING_MAX_PIXELS, SPOTTING_UPSCALE_THRESHOLD, ) from ocr_studio.spotting import TextSpan, parse_spans, spans_to_text, strip_special_tokens from ocr_studio.tiling import iter_tiles, join_tile_text, merge_tile_spans, offset_spans, should_tile TASK_OCR = "ocr" TASK_SPOTTING = "spotting" TASK_TABLE = "table" OFFICIAL_PROMPTS = { TASK_OCR: "OCR:", TASK_SPOTTING: "Spotting:", TASK_TABLE: "Table Recognition:", } LANGUAGE_HINTS = { LANGUAGE_FA: "The document is in Persian (Farsi). Preserve Persian letters, digits, and RTL order.\n", LANGUAGE_EN: "The document is in English.\n", LANGUAGE_BOTH: "The document mixes Persian (Farsi) and English. Preserve both scripts.\n", "auto": "The document may be Persian (Farsi), English, or mixed.\n", } TASK_INSTRUCTIONS = { MODE_DOCUMENT: "", MODE_PRECISE: "", MODE_TABLE: "Recover every table with cell structure.\n", MODE_MARKDOWN: "Transcribe as Markdown with headings, lists, and tables where they appear.\n", MODE_FIELDS: "Extract every labeled field as a line in the form 'label: value'. Include names, dates, amounts, IDs, and addresses.\n", MODE_COMPARE: "", } @dataclass class PageInference: raw_text: str display_text: str spans: list[TextSpan] task: str tile_count: int = 1 alt_text: str = "" @dataclass class BatchInference: pages: list[PageInference] = field(default_factory=list) compare: bool = False def build_prompt(mode: str, language_key: str, task: str) -> str: hint = LANGUAGE_HINTS.get(language_key, LANGUAGE_HINTS["auto"]) extra = TASK_INSTRUCTIONS.get(mode, "") token = OFFICIAL_PROMPTS.get(task, OFFICIAL_PROMPTS[TASK_OCR]) return f"{hint}{extra}{token}" def mode_to_task(mode: str) -> str: if mode == MODE_PRECISE: return TASK_SPOTTING if mode == MODE_TABLE: return TASK_TABLE return TASK_OCR class PaddleOcrVlEngine: def __init__(self) -> None: self.model: Any = None self.processor: Any = None self.device: Any = None self._lock = threading.Lock() def load(self) -> None: if self.model is not None: return from ocr_studio.config import MODEL_ID, MODEL_REVISION import torch from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor device = torch.device("cuda" if torch.cuda.is_available() else "cpu") dtype = torch.bfloat16 if device.type == "cuda" else torch.float32 config = AutoConfig.from_pretrained(MODEL_ID, revision=MODEL_REVISION) if not hasattr(config, "text_config") and hasattr(config, "get_text_config"): config.text_config = config.get_text_config() if getattr(config, "tie_word_embeddings", None): config.tie_word_embeddings = False text_config = getattr(config, "text_config", None) if text_config is not None and getattr(text_config, "tie_word_embeddings", None): text_config.tie_word_embeddings = False processor = AutoProcessor.from_pretrained( MODEL_ID, revision=MODEL_REVISION, trust_remote_code=False, ) model = AutoModelForImageTextToText.from_pretrained( MODEL_ID, config=config, revision=MODEL_REVISION, torch_dtype=dtype, trust_remote_code=False, low_cpu_mem_usage=True, ) model = model.to(device).eval() self.model = model self.processor = processor self.device = device def _prepare_image(self, image: Image.Image, task: str) -> Image.Image: prepared = image.convert("RGB") if ( task == TASK_SPOTTING and prepared.width < SPOTTING_UPSCALE_THRESHOLD and prepared.height < SPOTTING_UPSCALE_THRESHOLD ): prepared = prepared.resize( (prepared.width * 2, prepared.height * 2), Image.Resampling.LANCZOS, ) return prepared def _generate(self, image: Image.Image, prompt: str, task: str) -> str: self.load() import torch work_image = self._prepare_image(image, task) max_pixels = SPOTTING_MAX_PIXELS if task == TASK_SPOTTING else OCR_MAX_PIXELS messages = [ { "role": "user", "content": [ {"type": "image", "image": work_image}, {"type": "text", "text": prompt}, ], } ] image_processor = self.processor.image_processor min_pixels = getattr(image_processor, "min_pixels", None) if min_pixels is None: size_cfg = getattr(image_processor, "size", {}) or {} min_pixels = size_cfg.get("shortest_edge") or size_cfg.get("min_pixels") or (16 * 28 * 28) inputs = self.processor.apply_chat_template( messages, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt", images_kwargs={ "size": { "shortest_edge": int(min_pixels), "longest_edge": max_pixels, } }, ) inputs = inputs.to(self.model.device) with self._lock, torch.inference_mode(): generated = self.model.generate( **inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False, ) prompt_len = inputs["input_ids"].shape[-1] decoded = self.processor.decode( generated[0][prompt_len:], skip_special_tokens=True, clean_up_tokenization_spaces=False, ).strip() return decoded def _recognize_image( self, image: Image.Image, mode: str, language_key: str, task: str, use_tiles: bool, ) -> PageInference: prompt = build_prompt(mode, language_key, task) tiled = should_tile(image, use_tiles) and task != TASK_TABLE if not tiled: raw = self._generate(image, prompt, task) spans = parse_spans(raw, image.width, image.height) display = spans_to_text(spans, strip_special_tokens(raw)).strip() return PageInference(raw_text=raw, display_text=display, spans=spans, task=task, tile_count=1) texts: list[str] = [] spans: list[TextSpan] = [] tile_count = 0 for tile, origin_x, origin_y in iter_tiles(image): tile_count += 1 raw = self._generate(tile, prompt, task) tile_spans = offset_spans(parse_spans(raw, tile.width, tile.height), origin_x, origin_y) spans.extend(tile_spans) texts.append(spans_to_text(tile_spans, strip_special_tokens(raw)).strip()) spans = merge_tile_spans(spans) display = spans_to_text(spans, join_tile_text(texts)).strip() return PageInference( raw_text="\n".join(texts), display_text=display, spans=spans, task=task, tile_count=max(1, tile_count), ) def recognize_pages( self, pages: list[Image.Image], language_key: str, mode: str, high_accuracy: bool, ) -> BatchInference: self.load() results: list[PageInference] = [] compare = mode == MODE_COMPARE primary_task = TASK_SPOTTING if mode in {MODE_PRECISE, MODE_COMPARE} else mode_to_task(mode) use_tiles = high_accuracy and primary_task != TASK_TABLE for page in pages: primary = self._recognize_image(page, mode, language_key, primary_task, use_tiles) if compare: document = self._recognize_image(page, MODE_DOCUMENT, language_key, TASK_OCR, use_tiles) primary.alt_text = document.display_text if not primary.display_text: primary.display_text = document.display_text results.append(primary) return BatchInference(pages=results, compare=compare) def recognize(self, image: Image.Image, mode: str) -> PageInference: task = mode_to_task(mode) return self._recognize_image(image, mode, "auto", task, False)