| """Private document text extraction with OCR fallback for scanned PDFs.""" |
|
|
| from __future__ import annotations |
|
|
| import re |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
|
|
|
|
| @dataclass(frozen=True) |
| class ExtractedDocument: |
| text: str |
| method: str |
| pages: int |
| ocr_pages: int |
| truncated: bool |
|
|
| def public_dict(self) -> dict: |
| value = asdict(self) |
| value.pop("text", None) |
| value["text_chars"] = len(self.text) |
| return value |
|
|
|
|
| def _normalise(text: object) -> str: |
| value = str(text or "").replace("\x00", "") |
| value = re.sub(r"[ \t]+", " ", value) |
| value = re.sub(r"\n{3,}", "\n\n", value) |
| return value.strip() |
|
|
|
|
| def _trim(text: str, max_chars: int) -> tuple[str, bool]: |
| if len(text) <= max_chars: |
| return text, False |
| return text[:max_chars].rsplit(" ", 1)[0].rstrip() + "\n\n[Document truncated]", True |
|
|
|
|
| def extract_document(path: str | Path, media_type: str = "", *, max_chars: int = 120_000) -> ExtractedDocument: |
| file_path = Path(path) |
| suffix = file_path.suffix.lower() |
| if suffix in {".txt", ".md"}: |
| text, truncated = _trim(_normalise(file_path.read_text(encoding="utf-8-sig")), max_chars) |
| return ExtractedDocument(text, "text", 1, 0, truncated) |
| if suffix == ".docx": |
| from docx import Document |
|
|
| document = Document(str(file_path)) |
| blocks = [paragraph.text for paragraph in document.paragraphs if paragraph.text.strip()] |
| for table in document.tables: |
| for row in table.rows: |
| blocks.append(" | ".join(cell.text.strip() for cell in row.cells)) |
| text, truncated = _trim(_normalise("\n\n".join(blocks)), max_chars) |
| return ExtractedDocument(text, "docx", 1, 0, truncated) |
| if suffix != ".pdf" and media_type != "application/pdf": |
| raise ValueError("Unsupported document type for text extraction.") |
|
|
| import pymupdf as fitz |
|
|
| pages: list[str] = [] |
| ocr_pages = 0 |
| with fitz.open(str(file_path)) as pdf: |
| page_count = len(pdf) |
| for page in pdf: |
| value = _normalise(page.get_text("text")) |
| if len(value) < 40: |
| try: |
| import io |
| import pytesseract |
| from PIL import Image |
|
|
| pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False) |
| image = Image.open(io.BytesIO(pixmap.tobytes("png"))) |
| ocr = _normalise(pytesseract.image_to_string(image, lang="eng")) |
| if len(ocr) > len(value): |
| value = ocr |
| ocr_pages += 1 |
| except (ImportError, OSError, RuntimeError): |
| pass |
| if value: |
| pages.append(value) |
| text, truncated = _trim(_normalise("\n\n".join(pages)), max_chars) |
| method = "pdf+ocr" if ocr_pages else "pdf-text" |
| return ExtractedDocument(text, method, page_count, ocr_pages, truncated) |
|
|
