core/ocr_engine.py

"""
OCR utilities for the MarkItDown API.

Two engines are provided:

    ocr_image(source)
        RapidOCR singleton for raster images (JPEG, PNG, WEBP, etc.).
        Accepts bytes, a local file path string, an HTTP/HTTPS URL string,
        a numpy.ndarray, or a PIL.Image instance.

    ocr_pdf(source, dpi=150)
        Scanned-PDF fallback. Renders each page with pypdfium2, then feeds
        the rendered PIL image through ocr_image. Returns all pages joined
        with double newlines.

Both functions return a plain string and never raise; errors are logged and
an empty string is returned on failure.
"""

from __future__ import annotations

import io
import threading
from typing import Union
from urllib.parse import urlparse

import numpy as np

from logger import get_logger

logger = get_logger(__name__)


# ---------------------------------------------------------------------------
# RapidOCR singleton
# ---------------------------------------------------------------------------

_lock = threading.Lock()
_engine = None


def _get_engine():
    """Return the shared RapidOCR instance, initialising it on first call."""
    global _engine
    if _engine is None:
        with _lock:
            if _engine is None:
                from rapidocr_onnxruntime import RapidOCR
                _engine = RapidOCR(
                    Det={"use_cuda": False, "use_dml": False},
                    Cls={"use_cuda": False, "use_dml": False},
                    Rec={"use_cuda": False, "use_dml": False},
                    print_verbose=False,
                )
                logger.info("RapidOCR engine initialised")
    return _engine


# ---------------------------------------------------------------------------
# Input normalisation
# ---------------------------------------------------------------------------

def _to_numpy(source) -> Union[np.ndarray, str]:
    """Normalise *source* to a numpy array or a local file path string.

    Accepted input types:
        PIL.Image   — converted directly to ndarray.
        bytes       — decoded via PIL then converted to ndarray.
        str         — HTTP/HTTPS URL fetched then decoded; local paths returned as-is.
        np.ndarray  — returned unchanged.
    """
    from PIL import Image

    def _pil_to_array(img: Image.Image) -> np.ndarray:
        if img.mode not in ("RGB", "L", "RGBA"):
            img = img.convert("RGB")
        return np.array(img)

    if isinstance(source, np.ndarray):
        return source

    if isinstance(source, Image.Image):
        return _pil_to_array(source)

    if isinstance(source, (bytes, bytearray, memoryview)):
        return _pil_to_array(Image.open(io.BytesIO(bytes(source))))

    if isinstance(source, str):
        parsed = urlparse(source)
        if parsed.scheme in {"http", "https"}:
            import httpx
            resp = httpx.get(source, follow_redirects=True, timeout=30)
            resp.raise_for_status()
            return _pil_to_array(Image.open(io.BytesIO(resp.content)))
        # Local file path — RapidOCR accepts it directly.
        return source

    raise TypeError(
        f"ocr_image expects bytes, str (URL or path), numpy.ndarray, or PIL.Image; "
        f"received {type(source).__name__!r}"
    )


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def ocr_image(
    source,
    *,
    use_det: bool = True,
    use_cls: bool = True,
    use_rec: bool = True,
    text_score: float = 0.5,
) -> str:
    """Extract text from an image using RapidOCR.

    Parameters
    ----------
    source:
        Input image — bytes, URL string, local path string, numpy.ndarray,
        or PIL.Image.
    use_det, use_cls, use_rec:
        RapidOCR pipeline stages (detection, classification, recognition).
    text_score:
        Minimum confidence threshold for accepted text lines.

    Returns
    -------
    str
        Recognised text lines joined by newlines, or an empty string when
        no text is detected.
    """
    engine = _get_engine()
    img = _to_numpy(source)

    result, _ = engine(
        img,
        use_det=use_det,
        use_cls=use_cls,
        use_rec=use_rec,
        text_score=text_score,
    )

    if not result:
        return ""

    return "\n".join(item[1] for item in result if len(item) > 1 and item[1])


def ocr_pdf(source: Union[str, bytes], *, dpi: int = 150) -> str:
    """Extract text from a scanned (image-only) PDF using pypdfium2 and RapidOCR.

    Each page is rendered to a PIL image in memory (no temporary files are
    written), then passed through ocr_image. All page outputs are joined
    with double newlines.

    Parameters
    ----------
    source:
        Local file path (str) or raw PDF bytes.
    dpi:
        Rendering resolution. 150 balances speed and OCR quality for most
        document types. Increase to 200-300 for small or dense text.

    Returns
    -------
    str
        Concatenated OCR text from all pages, or an empty string on failure.
    """
    try:
        import pypdfium2 as pdfium
    except ImportError:
        logger.error("ocr_pdf | pypdfium2 not installed; run: pip install pypdfium2")
        return ""

    try:
        pdf = pdfium.PdfDocument(source)
        scale = dpi / 72.0  # pypdfium2 native resolution is 72 dpi
        page_texts: list[str] = []

        for page_index in range(len(pdf)):
            page = pdf[page_index]
            bitmap = page.render(scale=scale, rotation=0)
            pil_image = bitmap.to_pil()

            logger.debug("ocr_pdf | processing page %d/%d", page_index + 1, len(pdf))
            page_text = ocr_image(pil_image)
            if page_text:
                page_texts.append(page_text)

        pdf.close()
        return "\n\n".join(page_texts)

    except Exception as exc:
        logger.error("ocr_pdf | failed | error=%s", exc, exc_info=True)
        return ""