ocrcontext 0.1.0__py3-none-any.whl

ocrcontext/pipeline.py

@@ -0,0 +1,162 @@
+"""Document routing + OCR orchestration (no LLM here — that's the analyzer's job).
+
+Reproduces the retry/fallback ladder from app/api/documents/process/route.ts and
+the page loops in OCRService / HandwritingOCRService, minus the web/Modal layer.
+"""
+
+from __future__ import annotations
+
+from .config import AnalyzerConfig
+from .engines.base import OcrEngine, PageOcr
+from .engines.pdf_text import extract_pdf_text_preserve_layout, has_sufficient_pdf_text
+from .engines.registry import EngineRegistry
+from .quality import is_ocr_text_insufficient
+from .types import OcrResult, TextSource
+from .utils.files import (
+    Source,
+    cleanup_paths,
+    is_pdf,
+    load_source,
+    new_temp_path,
+    rasterize_pdf,
+)
+
+
+class Pipeline:
+    """Routes a document to the right engine(s) and returns raw OCR text."""
+
+    def __init__(
+        self,
+        registry: EngineRegistry | None = None,
+        config: AnalyzerConfig | None = None,
+    ) -> None:
+        self.registry = registry or EngineRegistry.shared()
+        self.config = config or AnalyzerConfig()
+
+    def run(
+        self,
+        source: Source,
+        *,
+        lang: str | None = None,
+        handwriting: bool = False,
+        filename: str | None = None,
+    ) -> OcrResult:
+        lang = lang or self.config.lang
+        file_bytes, ext = load_source(source, filename=filename)
+
+        # 1) Digital PDF text layer — exact text, no OCR / no GPU.
+        if is_pdf(ext) and self.config.prefer_pdf_text_layer and not handwriting:
+            full_text, page_count = extract_pdf_text_preserve_layout(file_bytes)
+            if has_sufficient_pdf_text(full_text):
+                return OcrResult(
+                    text=full_text.strip(),
+                    confidence=1.0,
+                    pages=page_count,
+                    text_source="pdf_text_layer",
+                )
+
+        # 2) Printed / scanned OCR (or handwriting if explicitly requested).
+        result = self._ocr(file_bytes, ext, lang=lang, handwriting=handwriting)
+
+        # 3) Auto handwriting fallback when printed OCR returns too little text.
+        if (
+            not handwriting
+            and self.config.auto_handwriting_fallback
+            and is_ocr_text_insufficient(result.text, result.pages)
+        ):
+            result = self._ocr(file_bytes, ext, lang=lang, handwriting=True)
+
+        return result
+
+    def _ocr(
+        self, file_bytes: bytes, ext: str, *, lang: str, handwriting: bool
+    ) -> OcrResult:
+        engine = (
+            self.registry.handwriting() if handwriting else self.registry.paddle()
+        )
+        render_scale = (
+            self.config.pdf_render_scale_handwriting
+            if handwriting
+            else self.config.pdf_render_scale
+        )
+        min_lines = (
+            self.config.min_lines_handwriting
+            if handwriting
+            else self.config.min_lines_per_page
+        )
+
+        image_paths, owned = self._materialize_images(file_bytes, ext, render_scale)
+        try:
+            return self._ocr_pages(
+                engine,
+                image_paths,
+                lang=lang,
+                handwriting=handwriting,
+                min_lines=min_lines,
+            )
+        finally:
+            cleanup_paths(owned)
+
+    @staticmethod
+    def _materialize_images(
+        file_bytes: bytes, ext: str, render_scale: float
+    ) -> tuple[list[str], list[str]]:
+        """Return (image_paths, owned_paths_to_cleanup)."""
+        if is_pdf(ext):
+            paths = rasterize_pdf(file_bytes, render_scale)
+            return paths, list(paths)
+        # Single image: write to a temp file the engines can read.
+        path = new_temp_path(ext)
+        with open(path, "wb") as f:
+            f.write(file_bytes)
+        return [path], [path]
+
+    @staticmethod
+    def _ocr_pages(
+        engine: OcrEngine,
+        image_paths: list[str],
+        *,
+        lang: str,
+        handwriting: bool,
+        min_lines: int,
+    ) -> OcrResult:
+        full_text = ""
+        all_scores: list[float] = []
+        used_vision = False
+        used_trocr = False
+        has_dikw = False
+
+        for idx, img_path in enumerate(image_paths):
+            if idx > 0:
+                full_text += f"\n\n--- Page {idx + 1} ---\n\n"
+
+            page: PageOcr = engine.recognize(
+                img_path, lang=lang, min_lines=min_lines, handwriting=handwriting
+            )
+            full_text += page.text + ("\n" if (handwriting and page.text) else "")
+            all_scores.extend(page.scores)
+            if page.has_dikw_structure:
+                has_dikw = True
+            if page.text_source == "vision_handwriting":
+                used_vision = True
+            elif page.text_source == "trocr_handwriting":
+                used_trocr = True
+
+        avg_conf = sum(all_scores) / len(all_scores) if all_scores else 0.0
+
+        text_source: TextSource = engine.text_source  # type: ignore[assignment]
+        if handwriting:
+            if used_vision and not used_trocr:
+                text_source = "vision_handwriting"
+            elif used_trocr:
+                text_source = "trocr_handwriting"
+            else:
+                text_source = "handwriting_ocr"
+
+        return OcrResult(
+            text=full_text.strip(),
+            confidence=round(avg_conf, 4),
+            pages=len(image_paths),
+            text_source=text_source,
+            has_dikw_structure=has_dikw,
+        )

ocrcontext/preprocessing/__init__.py

@@ -0,0 +1,5 @@
+"""Image preprocessing for OCR (deskew, denoise, CLAHE, line-band splitting)."""
+
+from .image import preprocess_image_for_ocr, split_image_into_line_bands
+
+__all__ = ["preprocess_image_for_ocr", "split_image_into_line_bands"]

ocrcontext/preprocessing/image.py

@@ -0,0 +1,177 @@
+"""Image preprocessing, ported verbatim from ocr-service/modal_app.py.
+
+OpenCV / NumPy are imported lazily so the base install does not require the
+``paddle`` extra unless an image OCR path is actually taken.
+"""
+
+from __future__ import annotations
+
+from ..exceptions import MissingDependencyError
+from ..utils.files import new_temp_path
+
+
+def _require_cv2():
+    """Import OpenCV, raising a friendly error if the image extras aren't installed."""
+    try:
+        import cv2
+
+        return cv2
+    except ImportError as exc:  # pragma: no cover - exercised via install matrix
+        raise MissingDependencyError("opencv-python-headless", "paddle") from exc
+
+
+def _deskew_grayscale(img):
+    """Correct slight rotation on notebook photos."""
+    import numpy as np
+
+    cv2 = _require_cv2()
+
+    inv = cv2.bitwise_not(img)
+    coords = np.column_stack(np.where(inv > 0))
+    if len(coords) < 200:
+        return img
+
+    angle = cv2.minAreaRect(coords)[-1]
+    if angle < -45:
+        angle = -(90 + angle)
+    else:
+        angle = -angle
+
+    if abs(angle) < 0.4 or abs(angle) > 12:
+        return img
+
+    h, w = img.shape[:2]
+    center = (w // 2, h // 2)
+    matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
+    return cv2.warpAffine(
+        img,
+        matrix,
+        (w, h),
+        flags=cv2.INTER_CUBIC,
+        borderMode=cv2.BORDER_REPLICATE,
+    )
+
+
+def _suppress_notebook_ruling(img):
+    """Reduce horizontal ruled lines that cross handwritten strokes."""
+    cv2 = _require_cv2()
+
+    h, w = img.shape[:2]
+    if h < 80 or w < 80:
+        return img
+
+    binary = cv2.adaptiveThreshold(
+        img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 15, 8
+    )
+    line_w = max(25, w // 25)
+    horiz_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (line_w, 1))
+    horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horiz_kernel, iterations=1)
+    horizontal = cv2.dilate(horizontal, horiz_kernel, iterations=1)
+
+    if cv2.countNonZero(horizontal) < 50:
+        return img
+
+    cleaned = cv2.inpaint(img, horizontal, 2, cv2.INPAINT_TELEA)
+    return cleaned
+
+
+def preprocess_image_for_ocr(img_path: str, handwriting: bool = False) -> str:
+    """Contrast + denoise before OCR (helps faint / handwritten scans).
+
+    Returns a path to a preprocessed PNG, or the original path if the image
+    could not be read. Caller is responsible for cleaning up new files.
+    """
+    cv2 = _require_cv2()
+
+    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+    if img is None:
+        return img_path
+
+    if handwriting:
+        h, w = img.shape[:2]
+        # Notebook margin rulers (e.g. 23, 22, 21...) confuse Vision - trim narrow left strip.
+        if w > h * 0.9 and w > 400:
+            crop_x = min(int(w * 0.08), 100)
+            img = img[:, crop_x:]
+        # Deskew only; skip ruled-line inpainting - it can erase strokes that touch notebook lines.
+        img = _deskew_grayscale(img)
+        img = cv2.fastNlMeansDenoising(img, None, h=8, templateWindowSize=7, searchWindowSize=21)
+        clahe = cv2.createCLAHE(clipLimit=2.5, tileGridSize=(8, 8))
+        img = clahe.apply(img)
+        # Slight upscale for thin strokes
+        img = cv2.resize(img, None, fx=1.15, fy=1.15, interpolation=cv2.INTER_CUBIC)
+    else:
+        clahe = cv2.createCLAHE(clipLimit=1.8, tileGridSize=(8, 8))
+        img = clahe.apply(img)
+
+    out_path = new_temp_path("png")
+    cv2.imwrite(out_path, img)
+    return out_path
+
+
+def split_image_into_line_bands(img_path: str) -> list[tuple[int, str]]:
+    """Create horizontal candidate bands for line recovery via projection profile.
+
+    Returns ``[(y_offset, band_image_path), ...]``. Caller owns cleanup of the
+    band image files. Ported from OCRService.process.split_image_into_line_bands.
+    """
+    import numpy as np
+
+    cv2 = _require_cv2()
+
+    bands: list[tuple[int, str]] = []
+    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+    if img is None:
+        return bands
+    h, w = img.shape[:2]
+    if h < 80 or w < 80:
+        return bands
+
+    blur = cv2.GaussianBlur(img, (3, 3), 0)
+    bw = cv2.adaptiveThreshold(
+        blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 31, 12
+    )
+
+    # Horizontal projection profile
+    row_sum = np.sum(bw > 0, axis=1)
+    threshold = max(8, int(0.02 * w))
+    active_rows = row_sum > threshold
+
+    segments: list[tuple[int, int]] = []
+    start = None
+    for i, active in enumerate(active_rows):
+        if active and start is None:
+            start = i
+        elif not active and start is not None:
+            if i - start >= 10:
+                segments.append((start, i))
+            start = None
+    if start is not None and (len(active_rows) - start) >= 10:
+        segments.append((start, len(active_rows)))
+
+    if not segments:
+        return bands
+
+    # Merge close segments (same line broken by weak strokes)
+    merged: list[list[int]] = []
+    for s, e in segments:
+        if not merged:
+            merged.append([s, e])
+        elif s - merged[-1][1] <= 8:
+            merged[-1][1] = e
+        else:
+            merged.append([s, e])
+
+    for idx_band, (s, e) in enumerate(merged):
+        pad = 10
+        y0 = max(0, s - pad)
+        y1 = min(h, e + pad)
+        crop = img[y0:y1, :]
+        if crop.shape[0] < 12:
+            continue
+        upscaled = cv2.resize(crop, None, fx=2.0, fy=2.0, interpolation=cv2.INTER_CUBIC)
+        band_path = new_temp_path("png")
+        cv2.imwrite(band_path, upscaled)
+        bands.append((y0, band_path))
+
+    return bands

ocrcontext/quality.py

@@ -0,0 +1,76 @@
+"""OCR text-quality heuristics, ported from lib/ocr/ocr-quality.ts and
+lib/ocr/handwriting-refine.ts / detect-dikw.ts.
+"""
+
+from __future__ import annotations
+
+import re
+
+from .types import RefinementMode
+
+_ALNUM_RE = re.compile(r"[^\W_]", re.UNICODE)  # unicode letters/digits
+_WS_RE = re.compile(r"\s+")
+
+
+def is_ocr_text_insufficient(text: str, page_count: int = 1) -> bool:
+    """Heuristic: OCR returned too little usable text (common on handwriting)."""
+    stripped = _WS_RE.sub(" ", text).strip()
+    if not stripped:
+        return True
+
+    pages = max(1, page_count)
+    min_chars = max(50, pages * 25)
+    if len(stripped) < min_chars:
+        return True
+
+    alnum = len(_ALNUM_RE.findall(stripped))
+    ratio = alnum / max(len(stripped), 1)
+    return ratio < 0.2
+
+
+# --- DIKW detection (ported from lib/ocr/detect-dikw.ts) -----------------------
+
+_DIKW_LETTERS = {"W", "K", "I", "D"}
+
+
+def _is_dikw_letter_token(token: str) -> bool:
+    t = re.sub(r"[^a-zA-Z]", "", token)
+    return len(t) == 1 and t.upper() in _DIKW_LETTERS
+
+
+def _is_pyramid_header(line: str) -> bool:
+    low = line.lower()
+    return "piramid" in low or "pyramid" in low or "dikw" in low
+
+
+def _row_looks_like_dikw_pair(line: str) -> bool:
+    tokens = [t for t in re.split(r"[\s·]+", line) if t]
+    if len(tokens) < 2:
+        return False
+    has_letter = any(_is_dikw_letter_token(t) for t in tokens)
+    has_long = any(len(re.sub(r"\W", "", t)) > 2 for t in tokens)
+    return has_letter and has_long
+
+
+def detect_dikw_structure(text: str) -> bool:
+    """True when OCR text looks like a DIKW / pyramid diagram (not plain prose)."""
+    lines = [ln.strip() for ln in text.split("\n") if ln.strip()]
+    pair_count = 0
+    for line in lines:
+        if _is_pyramid_header(line):
+            return True
+        if _row_looks_like_dikw_pair(line):
+            pair_count += 1
+            if pair_count >= 2:
+                return True
+    return False
+
+
+def handwriting_refinement_mode(
+    raw_text: str, has_dikw_from_ocr: bool | None = None
+) -> RefinementMode:
+    """Pick the handwriting refinement mode (layout for DIKW, else prose)."""
+    has_dikw = has_dikw_from_ocr is True or detect_dikw_structure(raw_text)
+    if has_dikw:
+        return RefinementMode.HANDWRITING_LAYOUT
+    return RefinementMode.HANDWRITING_PROSE

ocrcontext/schemas.py

@@ -0,0 +1,8 @@
+"""Public re-export of built-in extraction schemas.
+
+    from ocrcontext.schemas import Invoice, LineItem
+"""
+
+from .llm.schemas import INVOICE_EXTRACTION_PROMPT, Invoice, LineItem
+
+__all__ = ["Invoice", "LineItem", "INVOICE_EXTRACTION_PROMPT"]

ocrcontext/types.py

@@ -0,0 +1,55 @@
+"""Public result and value types."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class RefinementMode(str, Enum):
+    """LLM post-processing modes, ported verbatim from the original pipeline.
+
+    - ``layout``: digital PDFs — reconstruct clean structure.
+    - ``conservative``: printed OCR images/scans — minimal char-level correction.
+    - ``handwriting_layout``: handwritten notes/lists/tables/diagrams.
+    - ``handwriting_prose``: handwritten poems/paragraphs/letters.
+    """
+
+    LAYOUT = "layout"
+    CONSERVATIVE = "conservative"
+    HANDWRITING_LAYOUT = "handwriting_layout"
+    HANDWRITING_PROSE = "handwriting_prose"
+
+
+# String identifying which engine produced the text. Matches the original
+# `text_source` contract so downstream behaviour (e.g. skip-refine) is preserved.
+TextSource = Literal[
+    "pdf_text_layer",
+    "ocr",
+    "vision_handwriting",
+    "trocr_handwriting",
+    "handwriting_ocr",
+]
+
+
+class OcrResult(BaseModel):
+    """The output of an OCR / analysis run."""
+
+    text: str = Field(description="Extracted (and optionally refined) plain text.")
+    confidence: float = Field(default=0.0, description="Mean recognition confidence (0..1).")
+    pages: int = Field(default=1, description="Number of pages processed.")
+    text_source: TextSource = Field(
+        default="ocr", description="Which engine produced the text."
+    )
+    has_dikw_structure: bool = Field(
+        default=False, description="True when handwriting looks like a DIKW/pyramid diagram."
+    )
+    refined: bool = Field(default=False, description="True if an LLM refined the raw OCR text.")
+    raw_text: Optional[str] = Field(
+        default=None, description="Original OCR text before refinement (when refined)."
+    )
+
+    def __str__(self) -> str:  # convenient: print(result) -> the text
+        return self.text

ocrcontext/utils/__init__.py

@@ -0,0 +1 @@
+"""Internal utilities (language mapping, file handling)."""

ocrcontext/utils/files.py

@@ -0,0 +1,172 @@
+"""File / source loading and PDF rasterization helpers.
+
+Replaces the Modal service's ``/tmp`` plumbing with cross-platform temp handling
+(uses the OS temp dir via ``tempfile`` so it works on Windows too).
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import tempfile
+import uuid
+from contextlib import contextmanager
+from pathlib import Path
+from typing import IO, Iterator, Union
+
+from ..exceptions import UnsupportedFileError
+
+# What the public API accepts as a document source.
+Source = Union[str, Path, bytes, bytearray, IO[bytes]]
+
+IMAGE_EXTS = {"png", "jpg", "jpeg", "bmp", "tif", "tiff", "webp", "gif"}
+
+
+def is_ascii(text: str) -> bool:
+    return all(ord(ch) < 128 for ch in text)
+
+
+def short_path(path: str) -> str:
+    """Return the Windows 8.3 short path (ASCII) for an existing path.
+
+    PaddlePaddle/OpenCV's C++ file readers fail on paths containing non-ASCII
+    characters (a common case: a non-ASCII Windows username). The 8.3 short path
+    aliases the same file with ASCII-only characters. No-op off Windows or when
+    the path already resolves cleanly.
+    """
+    if sys.platform != "win32":
+        return path
+    try:
+        import ctypes
+        from ctypes import wintypes
+
+        _GetShortPathNameW = ctypes.windll.kernel32.GetShortPathNameW
+        _GetShortPathNameW.argtypes = [wintypes.LPCWSTR, wintypes.LPWSTR, wintypes.DWORD]
+        _GetShortPathNameW.restype = wintypes.DWORD
+
+        buf = ctypes.create_unicode_buffer(560)
+        result = _GetShortPathNameW(str(path), buf, len(buf))
+        if result:
+            return buf.value
+    except Exception:
+        pass
+    return path
+
+
+def ascii_safe_dir(path: str) -> str:
+    """ASCII-safe form of an existing directory (8.3 short path on Windows)."""
+    if is_ascii(path):
+        return path
+    candidate = short_path(path)
+    return candidate if is_ascii(candidate) else path
+
+
+def _temp_dir() -> str:
+    """ASCII-safe temp base. Override with ``OCRCONTEXT_TMPDIR``."""
+    base = os.environ.get("OCRCONTEXT_TMPDIR") or tempfile.gettempdir()
+    return ascii_safe_dir(base)
+
+
+def new_temp_path(ext: str) -> str:
+    """Absolute, ASCII-safe path to a unique temp file with the given extension."""
+    ext = ext.lstrip(".")
+    return os.path.join(_temp_dir(), f"ocrctx_{uuid.uuid4().hex}.{ext}")
+
+
+def load_source(source: Source, *, filename: str | None = None) -> tuple[bytes, str]:
+    """Normalize any accepted source into ``(file_bytes, extension)``.
+
+    ``extension`` is a lowercase string without a leading dot (e.g. ``"pdf"``).
+    """
+    # Path / path-like string
+    if isinstance(source, (str, Path)):
+        path = Path(source)
+        if not path.exists():
+            raise UnsupportedFileError(f"File not found: {path}")
+        ext = path.suffix.lower().lstrip(".")
+        if not ext:
+            raise UnsupportedFileError(f"Cannot determine file type from path: {path}")
+        return path.read_bytes(), ext
+
+    # Raw bytes — need a filename hint for the extension
+    if isinstance(source, (bytes, bytearray)):
+        ext = _ext_from_filename(filename)
+        return bytes(source), ext
+
+    # File-like object
+    if hasattr(source, "read"):
+        data = source.read()
+        if not isinstance(data, (bytes, bytearray)):
+            raise UnsupportedFileError("File-like source must be opened in binary mode.")
+        name = filename or getattr(source, "name", None)
+        ext = _ext_from_filename(name)
+        return bytes(data), ext
+
+    raise UnsupportedFileError(f"Unsupported source type: {type(source)!r}")
+
+
+def _ext_from_filename(filename: str | None) -> str:
+    if not filename or "." not in filename:
+        raise UnsupportedFileError(
+            "Could not infer file extension. Pass `filename=` when supplying raw bytes."
+        )
+    return filename.rsplit(".", 1)[-1].lower()
+
+
+def is_pdf(ext: str) -> bool:
+    return ext.lower() == "pdf"
+
+
+def is_image(ext: str) -> bool:
+    return ext.lower() in IMAGE_EXTS
+
+
+@contextmanager
+def temp_file(file_bytes: bytes, ext: str) -> Iterator[str]:
+    """Write bytes to a temp file, yield its path, and clean up afterwards."""
+    path = new_temp_path(ext)
+    with open(path, "wb") as f:
+        f.write(file_bytes)
+    try:
+        yield path
+    finally:
+        _safe_remove(path)
+
+
+def rasterize_pdf(file_bytes: bytes, scale: float, prefix: str = "page") -> list[str]:
+    """Render every PDF page to a PNG on disk; return the image paths.
+
+    Caller owns cleanup (use :func:`cleanup_paths`). Ported from the PDF render
+    loop in OCRService.process / HandwritingOCRService.process.
+    """
+    import fitz  # PyMuPDF — core dependency
+
+    from PIL import Image
+
+    image_paths: list[str] = []
+    pdf_document = fitz.open(stream=file_bytes, filetype="pdf")
+    try:
+        for page_num in range(len(pdf_document)):
+            page = pdf_document[page_num]
+            mat = fitz.Matrix(scale, scale)
+            pix = page.get_pixmap(matrix=mat)
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            out_path = new_temp_path("png")
+            img.save(out_path)
+            image_paths.append(out_path)
+    finally:
+        pdf_document.close()
+    return image_paths
+
+
+def cleanup_paths(paths: list[str]) -> None:
+    for p in paths:
+        _safe_remove(p)
+
+
+def _safe_remove(path: str) -> None:
+    try:
+        if path and os.path.exists(path):
+            os.remove(path)
+    except OSError:
+        pass