File size: 3,001 Bytes
1d9bd9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """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)
|