formula_img2latex/formula_img2latex/engines/p2t.py

"""Pix2Text formula OCR (MFR model) as a local candidate generator.

Not a vision-language model but a dedicated image-to-LaTeX transformer: another kind of model than
the VLM engines, so its errors are likely to be different ones. Runs in-process (ONNX), no server:
on the CPU, or on a GPU with onnxruntime-gpu (device "cuda:N"; N counts within CUDA_VISIBLE_DEVICES).
"""

from __future__ import annotations

import io
import re
import threading
import time

from PIL import Image

from ..errors import EXIT_TOOL, FormulaError
from .vlm import Recognition, prepare_image

DEFAULT_MODEL = "mfr-1.5"

# Pix2Text's bold: \mathrm{\boldmath~ x ~} (an error in math mode), \mathrm{{\bf~ r o t}}, \mathrm{\bf~ r o t}
_BOLDMATH = re.compile(r"\\mathrm\s*\{\s*\\boldmath\s*~?\s*([^{}~\\]*?)\s*~?\s*\}")
_BF = re.compile(r"\\mathrm\s*\{\s*(\{)?\s*\\bf\s*~?\s*([^{}~\\]*?)\s*~?\s*(?(1)\})\s*\}")


def tidy_p2t(tex: str) -> str:
    """Pix2Text output without superfluous spaces: `\\frac{\\alpha} {3 6 0}` -> `\\frac{\\alpha}{360}`.

    The spaces are removed by post_process of pix2tex (LaTeX-OCR), which keeps a space between two
    letters (`a x`, `\\cdot x`) and the control space `\\ `; TeX ignores the removed spaces in math
    mode, so the typeset formula does not change. Before that, Pix2Text's bold spellings become
    \\boldsymbol / \\mathbf, and its \\cfrac becomes \\dfrac (same look for ordinary fractions; Pandoc has
    no \\cfrac and would leave the formula as TeX text).
    """
    from pix2tex.utils.utils import post_process   # imports torch: only when needed

    tex = _BOLDMATH.sub(lambda m: "\\boldsymbol{" + re.sub(r"\s+", "", m.group(1)) + "}", tex)
    tex = _BF.sub(lambda m: "\\mathbf{" + re.sub(r"\s+", "", m.group(2)) + "}", tex)
    tex = re.sub(r"\\cfrac(?![A-Za-z])", r"\\dfrac", tex)
    return post_process(tex).strip()


def resolve_device(device: str) -> str:
    """Device "auto": the first visible GPU if onnxruntime has CUDA, else the CPU."""
    if device != "auto":
        return device
    import onnxruntime
    return "cuda" if "CUDAExecutionProvider" in onnxruntime.get_available_providers() else "cpu"


def ort_provider(device: str) -> tuple[str, dict]:
    """ONNX Runtime provider and its options for "cpu", "cuda" or "cuda:N".

    Set explicitly: pix2text otherwise takes the first available provider (TensorRT with onnxruntime-gpu).
    """
    if device == "cpu":
        return "CPUExecutionProvider", {}
    if device == "cuda" or device.startswith("cuda:"):
        return "CUDAExecutionProvider", {"device_id": int(device.partition(":")[2] or 0)}
    raise ValueError(f"unbekanntes Gerät {device!r} (cpu, cuda, cuda:N)")


class Pix2TextEngine:
    def __init__(self, model: str = DEFAULT_MODEL, device: str = "auto") -> None:
        self.name = "p2t"
        self.model_name = model
        self.model = f"pix2text-{model}"
        self.device = device
        self.thinking = None   # part of the cache key in pipeline.recognize; meaningless here
        self._ocr = None
        self._lock = threading.Lock()   # one model instance; inference is not shared between threads

    def check(self) -> None:
        """Load the model (downloads it to ~/.pix2text on first use)."""
        if self._ocr is not None:
            return
        try:
            from pix2text import LatexOCR
        except ImportError as e:
            raise FormulaError(f"pix2text ist nicht installiert ({e}): pip install 'pix2text[multilingual]'",
                               EXIT_TOOL) from e
        try:
            import pix2tex.utils.utils  # noqa: F401  (post_process for tidy_p2t)
        except ImportError as e:
            raise FormulaError(f"pix2tex ist nicht installiert ({e}): pip install --no-deps pix2tex munch",
                               EXIT_TOOL) from e
        self.device = resolve_device(self.device)
        provider, options = ort_provider(self.device)
        if provider == "CUDAExecutionProvider":
            import onnxruntime
            if provider not in onnxruntime.get_available_providers():
                raise FormulaError(f"Gerät {self.device}: onnxruntime ohne CUDA – pip install onnxruntime-gpu",
                                   EXIT_TOOL)
            onnxruntime.preload_dlls()   # cuDNN/cuBLAS from the nvidia-* pip packages (installed with torch)
        try:
            self._ocr = LatexOCR(model_name=self.model_name, model_backend="onnx", device=self.device,
                                 more_model_configs={"provider": provider, "provider_options": options})
        except Exception as e:   # download or model files broken: any exception type is possible here
            raise FormulaError(f"Pix2Text-Modell {self.model_name} nicht ladbar ({e})", EXIT_TOOL) from e

    def recognize(self, data: bytes, name: str = "", context: str = "", display: bool = False) -> Recognition:
        """Read one formula image. *context* and *display* are not used (the model sees only the image)."""
        t0 = time.time()
        try:
            self.check()
            img = Image.open(io.BytesIO(prepare_image(data, name))).convert("RGB")
            with self._lock:
                # _one_batch instead of recognize(): the latter draws a tqdm bar for every image
                out = self._ocr._one_batch([img], None)[0]
                raw = self._ocr.post_process(out["text"]).strip()
        except Exception as e:   # model not loadable or one unreadable image: must not stop the run
            return Recognition(self.name, self.model, None, None, "", str(e)[:300], time.time() - t0)
        score = float(out.get("score") or 0.0) or None
        tex = tidy_p2t(raw)
        return Recognition(self.name, self.model, tex or None, None, raw[:2000], None, time.time() - t0,
                           p_mean=score)