ocr_studio/engine.py

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)