diff --git a/.gitignore b/.gitignore index d5d9408..cd52464 100644 --- a/.gitignore +++ b/.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/ \ No newline at end of file diff --git a/OCR/README.md b/OCR/README.md new file mode 100644 index 0000000..13dd1f8 --- /dev/null +++ b/OCR/README.md @@ -0,0 +1,312 @@ +# OCR Text Extraction + +A standalone, reusable PDF text-extraction module with a **pluggable OCR engine architecture**. It converts PDF documents into clean, structured, machine-readable output — designed to be plugged into any downstream pipeline (search indexing, LLM ingestion, analytics, etc.) as an independent building block, either as a CLI tool or as an importable Python package. + +Ships today with [Surya OCR](https://github.com/datalab-to/surya) as the default (and only) built-in engine, but the OCR backend is fully swappable — see [Pluggable OCR engines](#pluggable-ocr-engines) below. + +## Why this exists + +Raw OCR output isn't directly usable — it's messy (per-block HTML fragments, layout labels, confidence scores mixed in), and it's usually tightly coupled to one specific OCR library. This project fixes both problems: + +- **Humans** get a readable, well-formatted HTML report to review documents visually. +- **Machines / other services** get plain text (simple to embed/search/feed to an LLM) **and** structured JSON (field-mapped, for programmatic consumption — databases, APIs, downstream processing). +- **The OCR backend itself is decoupled** behind a small interface, so swapping Surya for Tesseract, a vLLM-hosted vision model, a cloud OCR API, etc. never requires touching the pipeline, formatter, or CLI — only a new engine adapter. + +This makes text extraction a **decoupled, standalone concern**. Any other project can import this package directly (`from extractor import Extractor`) or just read the output files from a known folder, without depending on Surya, vllm/llama.cpp, or any OCR-specific code. + +## Scope (v1) + +- **Input**: PDFs only, for now. (Images like `.png`/`.jpg` are a future extension — the loader will be written so adding them later is trivial, but v1 only wires up PDFs.) +- **Input location**: defaults to `extraction_input/` (can contain PDFs directly, or PDFs nested inside subfolders), but is **not hardcoded** — the CLI accepts an optional path argument to point at a specific PDF or a different folder instead. See [CLI](#cli) below. +- **Output location**: one consolidated folder, `extraction_output/`, containing three subfolders — one per output format. Each processed PDF produces one file of the same base name in each subfolder. +- **OCR engine**: defaults to Surya, selectable via `--engine` (CLI) or the `engine=` constructor argument (package API). Only `surya` ships today, but the architecture supports adding more without changing any existing code. + +## Installation & Setup + +These steps are relative to this module's own folder (wherever it lives inside a larger repository) — they don't assume this is the repo root. + +### Prerequisites +- Python 3.11+ +- macOS/Linux (Surya's inference backend spawns a local `vllm` or `llama.cpp` process) + +### 1. Enter this module's folder +```bash +cd path/to/this/module # e.g. cd services/ocr-extraction +``` + +### 2. Create and activate a virtual environment (scoped to this module) +```bash +python3.11 -m venv surya-env +source surya-env/bin/activate # macOS/Linux +# surya-env\Scripts\activate # Windows +``` + +### 3. Install dependencies +```bash +pip install -r requirements.txt +``` +`requirements.txt` currently contains: +``` +surya-ocr +pypdfium2 +pillow +``` +`surya-ocr` pulls in Surya's inference stack (including its `vllm`/`llama.cpp` backend management) — no separate model download step is required; Surya downloads/spawns its own inference backend automatically on first use. + +### 4. Add your PDFs +Drop PDF files (optionally in nested subfolders) into `extraction_input/` inside this module's folder: +``` +extraction_input/ +├── some_document.pdf +└── nested_folder/ + └── another_document.pdf +``` + +### 5. Run the extraction +```bash +python main.py +``` +This processes every PDF under `extraction_input/` using the default `surya` engine, and writes `.html`, `.txt`, and `.json` outputs to `extraction_output/html/`, `extraction_output/text/`, `extraction_output/json/` respectively. + +**First run note:** Surya spawns its inference backend (vllm/llama.cpp) the first time OCR runs — this can take a bit longer on the very first invocation. Subsequent PDFs processed in the same `python main.py` run reuse the same backend instance (see [Use as a package](#use-as-a-package) for why). + +### 6. Check the output +```bash +cat extraction_output/text/some_document.txt # plain text +open extraction_output/html/some_document.html # human-readable report (macOS) +cat extraction_output/json/some_document.json # structured JSON +``` + +### Processing a single file or custom folder (without touching `extraction_input/`) +```bash +python main.py path/to/specific/file.pdf +python main.py path/to/some/other/folder +``` + +### Using it from another service in the same repo +If another service in this repo wants to call this module directly in-process rather than shelling out to `main.py`, see [Use as a package](#use-as-a-package) below — install this module's `requirements.txt` into that service's environment (or make this module pip-installable/a shared dependency), then `from extractor import Extractor`. + +## Project Structure + + +``` +OCR/ +├── extractor/ # Core, importable Python package +│ ├── __init__.py # Public API: exports `Extractor` (the only public entry point) +│ ├── models.py # Engine-agnostic result model: Block, PageResult dataclasses +│ ├── loader.py # PDF -> list of PIL Images (page-by-page rendering) +│ ├── formatter.py # Generic PageResult/Block -> plain text / HTML / JSON +│ ├── pipeline.py # `Extractor` class: orchestrates load -> OCR -> format -> save +│ └── engines/ # Pluggable OCR backends +│ ├── __init__.py # ENGINE_REGISTRY + get_engine() factory +│ ├── base.py # BaseOCREngine ABC — the contract every engine implements +│ └── surya_engine.py # SuryaEngine(BaseOCREngine) - wraps Surya's inference API +├── main.py # CLI entry point: `python main.py [path] [--engine NAME]` +├── extraction_input/ # Default input folder (PDFs, can have nested subfolders) +│ └── ... *.pdf +├── extraction_output/ # Hardcoded output folder (auto-created) +│ ├── html/ +│ │ └── .html +│ ├── text/ +│ │ └── .txt +│ └── json/ +│ └── .json +├── requirements.txt +└── README.md +``` + +## Pluggable OCR engines + +The pipeline never talks to Surya (or any OCR library) directly. It only depends on: + +1. **`extractor/engines/base.py` — `BaseOCREngine`**, an abstract class with a single required method: + ```python + class BaseOCREngine(ABC): + @abstractmethod + def run(self, images: List[PIL.Image]) -> List[PageResult]: + ... + ``` +2. **`extractor/models.py` — `Block` / `PageResult`**, a small engine-agnostic dataclass contract that every engine must translate its native output into: + ```python + @dataclass + class Block: + label: str # "Text", "SectionHeader", "Table", "Picture", ... + html: str # recognized content as HTML + bbox: List[float] # [x0, y0, x1, y1] + confidence: float = 1.0 + reading_order: int = 0 + skipped: bool = False # True for non-OCR'd blocks (e.g. pure images) + + @dataclass + class PageResult: + blocks: List[Block] + ``` + +`formatter.py` and `pipeline.py` are written entirely against this model — they have zero knowledge of Surya's actual object shapes. All Surya-specific translation logic lives in one place: `extractor/engines/surya_engine.py`. + +### Adding a new OCR engine + +1. Create `extractor/engines/your_engine.py`: + ```python + from .base import BaseOCREngine + from ..models import Block, PageResult + + class YourEngine(BaseOCREngine): + def run(self, images): + # call your OCR backend, then map its output into Block/PageResult + return [PageResult(blocks=[...]) for image in images] + ``` +2. Register it in `extractor/engines/__init__.py`: + ```python + ENGINE_REGISTRY = { + "surya": SuryaEngine, + "your_engine": YourEngine, + } + ``` +3. Done. Use it immediately, with **no other file needing changes**: + ```bash + python main.py --engine your_engine + ``` + ```python + Extractor(engine="your_engine") + ``` + +You can also skip the registry entirely and inject an already-constructed engine instance directly (handy for tests/mocks): +```python +Extractor(engine=YourEngine()) +``` + +## How it works (pipeline) + +1. **Discover** — `pipeline.py` recursively walks `extraction_input/` and collects every `.pdf` file (including nested subfolders). +2. **Load** — `loader.py` opens each PDF with `pypdfium2` and renders every page to a `PIL.Image` at a DPI/scale tuned for OCR accuracy. Handles cleanup (`pdf.close()`). +3. **OCR** — the selected engine (`extractor/engines/*.py`) runs all page images for a document through its OCR backend and returns a list of engine-agnostic `PageResult` objects (one per page). +4. **Format** — `formatter.py` converts that generic result into three parallel outputs (see below). +5. **Save** — each PDF's three outputs are written to `extraction_output/html/`, `extraction_output/text/`, `extraction_output/json/` using the PDF's base filename. + +## Output formats + +### 1. HTML (`extraction_output/html/.html`) — for humans +A styled, page-by-page report (clean typography, tables rendered as real ``, no layout-label clutter). Meant for visual review in a browser, not for parsing. + +### 2. Plain text (`extraction_output/text/.txt`) — for machines (simple case) +All block HTML stripped down to plain text, concatenated in reading order, with clear page-break markers. No markup, no metadata — just the text of the document, ready for embedding, search indexing, or feeding into an LLM prompt. This is the "lowest common denominator" format: any consumer can `open(...).read()` and get usable text with zero parsing logic. + +### 3. JSON (`extraction_output/json/.json`) — for machines (structured case) + +This needs a real field-mapping design, since OCR output is block-oriented (layout blocks with generic labels), not domain-specific fields. Here's how it maps into JSON: + +**What every engine gives us, per PDF (via the `Block`/`PageResult` model):** +- A list of pages +- Each page → list of blocks, each block has: + - `label` (canonicalized layout type: `Text`, `SectionHeader`, `PageHeader`, `Table`, `ListGroup`, `Picture`, ...) + - `html` (recognized content, as HTML — tables come back as full `
`, math as ``) + - `bbox` (position on the page) + - `confidence` (0–1, how sure the engine is) + - `reading_order` (0-indexed position in the page) + - `skipped` (true for pure visual blocks like `Picture`, not OCR'd) + +**How this maps into our JSON schema:** + +Rather than inventing arbitrary business fields (which OCR has no way of knowing — it doesn't know what a "case number" or "invoice total" is), the JSON output stays **structurally faithful to what the engine actually detected**, but cleaned up and renamed into a stable, predictable schema. Field mapping: + +| Our JSON field | Source (`Block`/`PageResult` field) | Notes | +|---|---|---| +| `document` | PDF filename | top-level identifier | +| `page_count` | number of pages processed | | +| `pages` | list, one entry per page | | +| `pages[i].page_number` | 1-indexed loop counter | | +| `pages[i].blocks` | `page_result.blocks` | filtered: `skipped` blocks excluded | +| `blocks[j].type` | `block.label` | renamed for clarity (`Text`, `Table`, `SectionHeader`, etc.) | +| `blocks[j].order` | `block.reading_order` | preserves reading order for reconstruction | +| `blocks[j].text` | `block.html`, HTML-stripped | plain text version of this block | +| `blocks[j].html` | `block.html` | raw HTML kept too — useful if a table/structure needs to be preserved (e.g. `
` markup for a table block) | +| `blocks[j].confidence` | `block.confidence` | lets downstream consumers filter/flag low-confidence blocks | +| `blocks[j].bbox` | `block.bbox` | `[x0, y0, x1, y1]` — useful if a consumer needs position (e.g. highlighting in a viewer) | + +This gives downstream consumers **three levels of granularity to choose from**, without us guessing at business-specific fields we can't reliably extract from generic OCR: +- Want just the text? → concatenate `blocks[*].text` in `order`. +- Want to reconstruct tables/structure? → use `blocks[*].html` where `type == "Table"`. +- Want to filter noisy content? → use `confidence` to drop/flag low-confidence blocks. +- Want visual mapping? → use `bbox`. + +**Important limitation to be upfront about:** OCR engines detect *layout* (what kind of block something is — text, header, table) but do **not** understand document semantics (they don't know "this text is the case number" or "this table's second column is the verdict"). So the JSON schema above is a **generic, faithful structuring of OCR output** — not a domain-specific extraction (like "case_number", "judge_name", etc.). If field-level business extraction is needed later (e.g. for legal documents — case number, parties, verdict), that would be a **separate downstream step** (e.g. regex/LLM-based extraction running on top of this JSON's `text`/`blocks`), layered on top of this project rather than baked into the OCR step itself. + +Example JSON shape: +```json +{ + "document": "S4", + "page_count": 4, + "pages": [ + { + "page_number": 1, + "blocks": [ + { + "type": "PageHeader", + "order": 0, + "text": "1 न्यायनिर्णय सं.पौ. ख. क्र. ७२२/२०२०", + "html": "

...

", + "confidence": 0.97, + "bbox": [34.0, 12.0, 560.0, 48.0] + }, + { + "type": "Text", + "order": 1, + "text": "प्राप्त दिनांक : २७.०७.२०२०...", + "html": "

...

", + "confidence": 0.95, + "bbox": [34.0, 60.0, 560.0, 140.0] + } + ] + } + ] +} +``` + +## Use as a package + +`Extractor` is the **single public entry point** of this package — everything (CLI included) goes through it. + +```python +from extractor import Extractor + +# Pick an engine by name (default is "surya") +extractor = Extractor(engine="surya") + +# Process one PDF in-memory (no files written) - great for embedding in another service +result = extractor.process_pdf("extraction_input/some/nested/file.pdf") + +result.text # plain text string +result.json_data # dict, same shape as the JSON file +result.html # HTML report string + +# Or process a whole folder and write html/text/json outputs to disk +results = extractor.run(input_path="extraction_input/", output_dir="extraction_output/") +``` + +You can also inject a custom or mocked engine instance directly, instead of a name: +```python +from extractor import Extractor +from my_custom_engine import MyEngine + +extractor = Extractor(engine=MyEngine()) +``` + +Since the same `Extractor` object reuses one engine instance across every PDF it processes, any backend server the engine spawns (e.g. Surya's vllm/llama.cpp process) only starts once per `Extractor`, not once per document — construct it once and reuse it for batch runs. + +## CLI + +```bash +python main.py # processes everything under extraction_input/ (default) +python main.py path/to/file.pdf # processes a single specific PDF +python main.py path/to/folder # recursively processes all PDFs in a custom folder +python main.py path/to/folder --engine surya # explicitly select an OCR engine (default: surya) +``` +Recursively finds every `.pdf` under the given path (or `extraction_input/` if no path is given), runs OCR, and writes matching `.html`, `.txt`, `.json` files into `extraction_output/html/`, `extraction_output/text/`, `extraction_output/json/` respectively (same base filename as the source PDF, folder structure flattened by filename — if there's a name collision across nested folders, the relative path will be used to disambiguate). + +## Accuracy considerations (multilingual documents) + +Since sample documents are in Hindi/Marathi/mixed scripts: +- Render scale defaults to produce ~150–200 DPI equivalent (capped at ~2048px width) for legible text — this is the single biggest accuracy lever per Surya's own docs. +- Room to add preprocessing (deskew/binarize) for scanned/blurry documents later, without changing the public API. +- `confidence` scores are surfaced in JSON specifically so low-quality OCR on non-Latin scripts can be flagged/reviewed rather than silently trusted. diff --git a/OCR/extraction_input/.DS_Store b/OCR/extraction_input/.DS_Store new file mode 100644 index 0000000..2396882 Binary files /dev/null and b/OCR/extraction_input/.DS_Store differ diff --git a/OCR/extraction_input/image_table.pdf b/OCR/extraction_input/image_table.pdf new file mode 100644 index 0000000..c42a940 Binary files /dev/null and b/OCR/extraction_input/image_table.pdf differ diff --git a/OCR/extractor/__init__.py b/OCR/extractor/__init__.py new file mode 100644 index 0000000..31b4d13 --- /dev/null +++ b/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", +] diff --git a/OCR/extractor/engines/__init__.py b/OCR/extractor/engines/__init__.py new file mode 100644 index 0000000..fbd5189 --- /dev/null +++ b/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"] diff --git a/OCR/extractor/engines/base.py b/OCR/extractor/engines/base.py new file mode 100644 index 0000000..8743fa7 --- /dev/null +++ b/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 diff --git a/OCR/extractor/engines/surya_engine.py b/OCR/extractor/engines/surya_engine.py new file mode 100644 index 0000000..2e53db6 --- /dev/null +++ b/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) diff --git a/OCR/extractor/formatter.py b/OCR/extractor/formatter.py new file mode 100644 index 0000000..24c7eb6 --- /dev/null +++ b/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"", "\n", html) + text = re.sub(r"", "\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 = [ + "", + f"{document_name}", + "", + ] + + for page_number, page_result in enumerate(pages, start=1): + parts.append("
") + parts.append(f"
Page {page_number}
") + for block in page_result.blocks: + if block.skipped: + continue + parts.append("
") + parts.append(block.html) + parts.append("
") + parts.append("
") + + parts.append("") + return "\n".join(parts) diff --git a/OCR/extractor/loader.py b/OCR/extractor/loader.py new file mode 100644 index 0000000..69d6946 --- /dev/null +++ b/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}") diff --git a/OCR/extractor/models.py b/OCR/extractor/models.py new file mode 100644 index 0000000..b5de133 --- /dev/null +++ b/OCR/extractor/models.py @@ -0,0 +1,48 @@ +"""Engine-agnostic data model shared by every OCR engine and by `formatter.py`. + +Any OCR engine (Surya, Tesseract, a vLLM-backed model, a cloud OCR API, ...) +must translate its own native output into these two dataclasses inside its +`run()` method. Nothing downstream (formatter, pipeline) ever imports or +touches an engine-specific type - this is the single contract that makes +engines swappable. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List + + +@dataclass +class Block: + """A single recognized content block on a page (a paragraph, heading, + table, list item, etc.) + + Attributes: + label: canonical block type, e.g. "Text", "SectionHeader", "Table", + "PageHeader", "Picture", "ListGroup". + html: recognized content as HTML (tables as `
`, plain + paragraphs as `

`, etc.) + bbox: `[x0, y0, x1, y1]` position of the block on the rendered page. + confidence: 0-1 engine confidence score for this block. Engines that + don't provide one should default to `1.0`. + reading_order: 0-indexed position of this block within the page, in + reading order. + skipped: True for blocks that were detected but not OCR'd (e.g. pure + images/figures). Skipped blocks are excluded from text/JSON/HTML + output by `formatter.py`. + """ + + label: str + html: str + bbox: List[float] + confidence: float = 1.0 + reading_order: int = 0 + skipped: bool = False + + +@dataclass +class PageResult: + """All recognized blocks for a single page, in reading order.""" + + blocks: List[Block] = field(default_factory=list) diff --git a/OCR/extractor/pipeline.py b/OCR/extractor/pipeline.py new file mode 100644 index 0000000..f8f8215 --- /dev/null +++ b/OCR/extractor/pipeline.py @@ -0,0 +1,136 @@ +"""High-level orchestration: find PDFs -> load -> OCR -> format -> save. + +Public API: + from extractor import Extractor + + extractor = Extractor(engine="surya") # or engine= + result = extractor.process_pdf("file.pdf") # single PDF, in-memory, no disk writes + results = extractor.run(input_path="folder/") # full folder pipeline, writes outputs to disk + +`Extractor` is intentionally the *only* public surface of this package - +whether you're running the CLI (`main.py`) or importing this as a library +into another service, you go through this one class. It depends only on +the abstract `BaseOCREngine` interface (see `engines/base.py`) and the +engine-agnostic `models.PageResult`/`models.Block` data (see `models.py`), +so swapping the OCR backend never requires touching this file. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import List, Optional, Union + +from . import formatter +from .engines import BaseOCREngine, get_engine +from .loader import find_pdfs, load_pages + +DEFAULT_INPUT_DIR = "extraction_input" +DEFAULT_OUTPUT_DIR = "extraction_output" +DEFAULT_ENGINE = "surya" + + +@dataclass +class ExtractionResult: + document: str # base filename, no extension + source_path: Path + text: str + json_data: dict + html: str + + +class Extractor: + """Runs the full PDF -> OCR -> formatted-output pipeline. + + Args: + engine: either a registered engine name (e.g. "surya"), or an + already-constructed `BaseOCREngine` instance (useful for + injecting a custom or mocked engine). Defaults to "surya". + The resolved engine is instantiated once and reused across + every PDF processed by this `Extractor` instance, so any + backend server it spawns (e.g. vllm/llama.cpp) only starts once. + """ + + def __init__(self, engine: Union[str, BaseOCREngine] = DEFAULT_ENGINE) -> None: + self.engine: BaseOCREngine = get_engine(engine) + + def process_pdf(self, pdf_path: Union[str, Path]) -> ExtractionResult: + """Run the full extraction pipeline on a single PDF and return the + result. Does not write any files - use `run()` for that, or save + the returned result's fields yourself. + """ + pdf_path = Path(pdf_path) + + images = load_pages(pdf_path) + pages = self.engine.run(images) + + document_name = pdf_path.stem + return ExtractionResult( + document=document_name, + source_path=pdf_path, + text=formatter.build_text(pages), + json_data=formatter.build_json(document_name, pages), + html=formatter.build_html(document_name, pages), + ) + + def run( + self, + input_path: Union[str, Path] = DEFAULT_INPUT_DIR, + output_dir: Union[str, Path] = DEFAULT_OUTPUT_DIR, + ) -> List[ExtractionResult]: + """Find every PDF under `input_path`, OCR it, and write outputs to + `output_dir/{html,text,json}/`. + + Returns the list of `ExtractionResult` objects (also useful if you + want to use the results in-process without re-reading the files + back). + """ + input_path = Path(input_path) + output_dir = Path(output_dir) + + html_dir = output_dir / "html" + text_dir = output_dir / "text" + json_dir = output_dir / "json" + for d in (html_dir, text_dir, json_dir): + d.mkdir(parents=True, exist_ok=True) + + pdfs = find_pdfs(input_path) + if not pdfs: + print(f"No PDF files found under: {input_path}") + return [] + + root = input_path if input_path.is_dir() else input_path.parent + + results = [] + for pdf_path in pdfs: + print(f"Processing: {pdf_path}") + result = self.process_pdf(pdf_path) + out_name = self._output_name(pdf_path, root) + + (text_dir / f"{out_name}.txt").write_text(result.text, encoding="utf-8") + (html_dir / f"{out_name}.html").write_text(result.html, encoding="utf-8") + (json_dir / f"{out_name}.json").write_text( + json.dumps(result.json_data, ensure_ascii=False, indent=2), + encoding="utf-8", + ) + + print(f" -> {text_dir / f'{out_name}.txt'}") + print(f" -> {html_dir / f'{out_name}.html'}") + print(f" -> {json_dir / f'{out_name}.json'}") + + results.append(result) + + return results + + @staticmethod + def _output_name(pdf_path: Path, root: Path) -> str: + """Derive a unique output base-name for a PDF, disambiguating + collisions across nested folders by including the relative path. + """ + try: + rel = pdf_path.relative_to(root).with_suffix("") + return str(rel).replace("/", "__") + except ValueError: + # pdf_path wasn't under root (e.g. a single file was passed directly) + return pdf_path.stem diff --git a/OCR/main.py b/OCR/main.py new file mode 100644 index 0000000..0552c58 --- /dev/null +++ b/OCR/main.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +"""Command-line entry point for the OCR text-extraction pipeline. + +Usage: + python main.py # processes everything under extraction_input/ + python main.py path/to/file.pdf # processes a single specific PDF + python main.py path/to/folder # recursively processes all PDFs in a custom folder + python main.py path/to/folder --engine surya # explicitly select an OCR engine (default: surya) +""" + +import argparse + +from extractor import DEFAULT_ENGINE, DEFAULT_INPUT_DIR, Extractor +from extractor.engines import ENGINE_REGISTRY + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run OCR text extraction on PDFs.") + parser.add_argument( + "input_path", + nargs="?", + default=DEFAULT_INPUT_DIR, + help=f"PDF file or folder to process (default: {DEFAULT_INPUT_DIR}/)", + ) + parser.add_argument( + "--engine", + default=DEFAULT_ENGINE, + choices=sorted(ENGINE_REGISTRY), + help=f"OCR engine to use (default: {DEFAULT_ENGINE})", + ) + + args = parser.parse_args() + + extractor = Extractor(engine=args.engine) + results = extractor.run(input_path=args.input_path) + + print(f"\nDone. Processed {len(results)} PDF(s).") + + +if __name__ == "__main__": + main() diff --git a/OCR/requirements.txt b/OCR/requirements.txt new file mode 100644 index 0000000..a72f30a --- /dev/null +++ b/OCR/requirements.txt @@ -0,0 +1,3 @@ +surya-ocr +pypdfium2 +pillow