Add ocr service

.gitignore

@@ -9,6 +9,7 @@ legal_cases_export.csv
 court_case_pipeline_v3.ipynb
 detailed_gap_analysis.md
 .legalai/
+.DS_store
 
 # Database storage
 qdrant_storage/

OCR/extractor/__init__.py

@@ -0,0 +1,24 @@
+"""OCR-powered PDF text-extraction package.
+
+Public API:
+    from extractor import Extractor
+
+    extractor = Extractor(engine="surya")           # pick/plug in any registered engine
+    result  = extractor.process_pdf("file.pdf")      # single PDF, in-memory
+    results = extractor.run(input_path="folder/")    # full folder pipeline, writes to disk
+
+`Extractor` is the sole public entry point of this package - see
+`extractor/pipeline.py` for details, `extractor/engines/base.py` for the
+OCR engine contract, and `extractor/models.py` for the engine-agnostic
+result data model.
+"""
+
+from .pipeline import DEFAULT_ENGINE, DEFAULT_INPUT_DIR, DEFAULT_OUTPUT_DIR, Extractor, ExtractionResult
+
+__all__ = [
+    "Extractor",
+    "ExtractionResult",
+    "DEFAULT_ENGINE",
+    "DEFAULT_INPUT_DIR",
+    "DEFAULT_OUTPUT_DIR",
+]

OCR/extractor/engines/__init__.py

@@ -0,0 +1,54 @@
+"""OCR engine registry.
+
+Adding a new engine is a two-step process:
+    1. Implement it in a new module here (subclassing `BaseOCREngine`).
+    2. Add one line to `ENGINE_REGISTRY` below.
+
+Nothing else in the codebase needs to change - `pipeline.py` only ever
+resolves engines through `get_engine()`.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, Type, Union
+
+from .base import BaseOCREngine
+from .surya_engine import SuryaEngine
+
+ENGINE_REGISTRY: Dict[str, Type[BaseOCREngine]] = {
+    "surya": SuryaEngine,
+}
+
+
+def get_engine(engine: Union[str, BaseOCREngine] = "surya") -> BaseOCREngine:
+    """Resolve an engine name (or an already-instantiated engine) into a
+    ready-to-use `BaseOCREngine` instance.
+
+    Args:
+        engine: either a registered engine name (e.g. "surya"), or an
+            already-constructed `BaseOCREngine` instance (passed through
+            unchanged, so callers can inject a custom/mocked engine).
+
+    Returns:
+        A `BaseOCREngine` instance.
+    """
+    if isinstance(engine, BaseOCREngine):
+        return engine
+
+    if isinstance(engine, str):
+        try:
+            engine_cls = ENGINE_REGISTRY[engine]
+        except KeyError as exc:
+            available = ", ".join(sorted(ENGINE_REGISTRY))
+            raise ValueError(
+                f"Unknown OCR engine '{engine}'. Available engines: {available}"
+            ) from exc
+        return engine_cls()
+
+    raise TypeError(
+        "engine must be a registered engine name (str) or a BaseOCREngine "
+        f"instance, got {type(engine)!r}"
+    )
+
+
+__all__ = ["BaseOCREngine", "ENGINE_REGISTRY", "get_engine"]

OCR/extractor/engines/base.py

@@ -0,0 +1,41 @@
+"""Abstract base class every OCR engine must implement.
+
+This is the single seam that makes OCR backends swappable. `pipeline.py`
+and `formatter.py` only ever talk to this interface (and the `models.py`
+data model it returns) - they never know or care whether the concrete
+implementation is Surya, Tesseract, a vLLM-hosted vision model, a cloud OCR
+API, etc.
+
+To add a new engine:
+    1. Create `extractor/engines/your_engine.py`.
+    2. Subclass `BaseOCREngine` and implement `run()`, translating your
+       engine's native output into `models.Block` / `models.PageResult`.
+    3. Register it in `extractor/engines/__init__.py`'s `ENGINE_REGISTRY`.
+
+Nothing else in the codebase needs to change.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import List
+
+from PIL import Image
+
+from ..models import PageResult
+
+
+class BaseOCREngine(ABC):
+    """Contract for any OCR backend used by the extraction pipeline."""
+
+    @abstractmethod
+    def run(self, images: List[Image.Image]) -> List[PageResult]:
+        """Run OCR on a list of page images for a single document.
+
+        Args:
+            images: one `PIL.Image` per page, in page order.
+
+        Returns:
+            One `PageResult` per input image, in the same order.
+        """
+        raise NotImplementedError

OCR/extractor/engines/surya_engine.py

@@ -0,0 +1,59 @@
+"""Surya OCR engine implementation.
+
+Thin adapter around Surya OCR's inference manager + recognition predictor.
+All Surya-specific knowledge (its API shape, its raw `PageOCRResult`/`block`
+attributes) is contained entirely in this file - translated into the
+engine-agnostic `models.Block` / `models.PageResult` before leaving `run()`.
+If Surya's API changes again (as it did between v1 and v2), only this file
+needs to change.
+"""
+
+from __future__ import annotations
+
+from typing import List
+
+from PIL import Image
+from surya.inference import SuryaInferenceManager
+from surya.recognition import RecognitionPredictor
+
+from ..models import Block, PageResult
+from .base import BaseOCREngine
+
+
+class SuryaEngine(BaseOCREngine):
+    """Lazily spins up the Surya inference backend and reuses it across calls.
+
+    Instantiate this once (e.g. via `Extractor`) and reuse it for every PDF
+    in a batch run, so the underlying vllm/llama.cpp server is only spawned
+    once instead of once per document.
+    """
+
+    def __init__(self) -> None:
+        self._manager = None
+        self._recognizer = None
+
+    def _ensure_ready(self) -> None:
+        if self._manager is None:
+            self._manager = SuryaInferenceManager()  # auto-spawns vllm or llama-server
+            self._recognizer = RecognitionPredictor(self._manager)
+
+    def run(self, images: List[Image.Image]) -> List[PageResult]:
+        self._ensure_ready()
+        raw_predictions = self._recognizer(images)
+        return [self._to_page_result(page) for page in raw_predictions]
+
+    @staticmethod
+    def _to_page_result(raw_page) -> PageResult:
+        """Translate a Surya `PageOCRResult` into our generic `PageResult`."""
+        blocks = [
+            Block(
+                label=block.label,
+                html=block.html,
+                bbox=list(block.bbox),
+                confidence=block.confidence,
+                reading_order=block.reading_order,
+                skipped=block.skipped,
+            )
+            for block in raw_page.blocks
+        ]
+        return PageResult(blocks=blocks)

OCR/extractor/formatter.py

@@ -0,0 +1,125 @@
+"""Converts engine-agnostic OCR output (`models.PageResult`/`models.Block`)
+into three consumer-facing formats:
+
+- plain text  (machines, simple case)
+- JSON        (machines, structured case)
+- HTML        (humans)
+
+Operates purely against the generic data model in `models.py` - it has no
+knowledge of which OCR engine produced the results, so swapping engines
+never requires touching this file.
+
+See README.md "Output formats" section for the full design rationale and
+field-mapping explanation.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Dict, List
+
+from .models import PageResult
+
+_TAG_RE = re.compile(r"<[^>]+>")
+
+
+def html_to_text(html: str) -> str:
+    """Strip HTML tags down to plain text.
+
+    Good enough for OCR output, which is simple HTML (p/b/u/table/tr/td/br),
+    not full documents - we just need readable text, not a full HTML parser.
+    """
+    # Turn common block/line-break tags into newlines/spaces before stripping,
+    # so table cells and paragraphs don't get smashed together.
+    text = re.sub(r"<br\s*/?>", "\n", html)
+    text = re.sub(r"</(p|tr|div|li)>", "\n", text)
+    text = re.sub(r"<(td|th)>", " ", text)
+    text = _TAG_RE.sub("", text)
+    # Collapse excess whitespace but keep paragraph breaks.
+    lines = [line.strip() for line in text.splitlines()]
+    lines = [line for line in lines if line]
+    return "\n".join(lines)
+
+
+def build_json(document_name: str, pages: List[PageResult]) -> Dict[str, Any]:
+    """Build the structured JSON representation for one document.
+
+    See README.md for the field mapping table (engine field -> our field).
+    """
+    page_entries = []
+    for page_number, page_result in enumerate(pages, start=1):
+        blocks = []
+        for block in page_result.blocks:
+            if block.skipped:
+                continue
+            blocks.append(
+                {
+                    "type": block.label,
+                    "order": block.reading_order,
+                    "text": html_to_text(block.html),
+                    "html": block.html,
+                    "confidence": block.confidence,
+                    "bbox": list(block.bbox),
+                }
+            )
+        page_entries.append({"page_number": page_number, "blocks": blocks})
+
+    return {
+        "document": document_name,
+        "page_count": len(pages),
+        "pages": page_entries,
+    }
+
+
+def build_text(pages: List[PageResult]) -> str:
+    """Build the plain-text representation for one document.
+
+    Concatenates every non-skipped block's text, in reading order, with
+    clear page-break markers. No markup, no metadata.
+    """
+    page_texts = []
+    for page_number, page_result in enumerate(pages, start=1):
+        block_texts = [
+            html_to_text(block.html)
+            for block in page_result.blocks
+            if not block.skipped
+        ]
+        page_body = "\n\n".join(block_texts)
+        page_texts.append(f"--- Page {page_number} ---\n{page_body}")
+
+    return "\n\n".join(page_texts)
+
+
+_HTML_STYLE = """
+body { font-family: sans-serif; padding: 20px; }
+.page { border: 1px solid #ccc; margin-bottom: 30px; padding: 15px; border-radius: 8px; }
+.page-title { background: #333; color: #fff; padding: 6px 12px; margin: -15px -15px 15px -15px;
+border-radius: 8px 8px 0 0; font-weight: bold; }
+.block { margin-bottom: 12px; }
+table, td, th { border: 1px solid #999; border-collapse: collapse; padding: 4px; }
+"""
+
+
+def build_html(document_name: str, pages: List[PageResult]) -> str:
+    """Build the styled, human-readable HTML report for one document."""
+    parts = [
+        "<html><head><meta charset='utf-8'>",
+        f"<title>{document_name}</title>",
+        "<style>",
+        _HTML_STYLE,
+        "</style></head><body>",
+    ]
+
+    for page_number, page_result in enumerate(pages, start=1):
+        parts.append("<div class='page'>")
+        parts.append(f"<div class='page-title'>Page {page_number}</div>")
+        for block in page_result.blocks:
+            if block.skipped:
+                continue
+            parts.append("<div class='block'>")
+            parts.append(block.html)
+            parts.append("</div>")
+        parts.append("</div>")
+
+    parts.append("</body></html>")
+    return "\n".join(parts)

OCR/extractor/loader.py

@@ -0,0 +1,64 @@
+"""Loading utilities: turn a PDF file into a list of page images.
+
+Kept deliberately narrow (PDF-only for v1) but structured so that adding
+support for plain image files (.png/.jpg/etc.) later is a small addition -
+just another branch in `load_pages`.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List
+
+import pypdfium2 as pdfium
+from PIL import Image
+
+# Render scale tuned for OCR accuracy on Indic/mixed-script documents.
+# scale=2 ~= 144 DPI (pdfium's base unit is 72 DPI). Surya's own docs
+# recommend keeping images under ~2048px width for the best
+# accuracy/throughput tradeoff - scale=2 keeps typical A4 pages comfortably
+# under that while still being sharp enough for OCR.
+DEFAULT_RENDER_SCALE = 2.0
+
+
+def load_pages(pdf_path: str | Path, scale: float = DEFAULT_RENDER_SCALE) -> List[Image.Image]:
+    """Render every page of a PDF to a PIL Image.
+
+    Args:
+        pdf_path: path to a .pdf file.
+        scale: pdfium render scale (2.0 ~= 144 DPI).
+
+    Returns:
+        List of PIL Images, one per page, in page order.
+    """
+    pdf_path = Path(pdf_path)
+    if pdf_path.suffix.lower() != ".pdf":
+        raise ValueError(f"Only PDF files are supported right now, got: {pdf_path}")
+
+    pdf = pdfium.PdfDocument(str(pdf_path))
+    try:
+        images = [page.render(scale=scale).to_pil() for page in pdf]
+    finally:
+        pdf.close()  # always release pdfium's native handles
+
+    return images
+
+
+def find_pdfs(input_path: str | Path) -> List[Path]:
+    """Resolve a CLI input path into a list of PDF files to process.
+
+    - If `input_path` is a single .pdf file, returns just that file.
+    - If `input_path` is a directory, recursively finds every .pdf under it
+      (including nested subfolders).
+    """
+    input_path = Path(input_path)
+
+    if input_path.is_file():
+        if input_path.suffix.lower() != ".pdf":
+            raise ValueError(f"Not a PDF file: {input_path}")
+        return [input_path]
+
+    if input_path.is_dir():
+        return sorted(p for p in input_path.rglob("*.pdf"))
+
+    raise FileNotFoundError(f"Input path does not exist: {input_path}")