88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
|
"""
|
||
|
|
OCR scanned PDFs using Tesseract + PyMuPDF (no poppler needed).
|
||
|
|
Reads each unparseable PDF page-by-page, renders to image, OCRs the image,
|
||
|
|
and saves the text to data/review/ocr/<invoice_no>.txt.
|
||
|
|
|
||
|
|
Then merges with the text-based extraction by writing OCR'd text back into
|
||
|
|
a fake pdfplumber-compatible page object via monkey-patching.
|
||
|
|
|
||
|
|
Run: python scripts/ocr_invoices.py
|
||
|
|
"""
|
||
|
|
import os, sys, subprocess, json
|
||
|
|
from pathlib import Path
|
||
|
|
import fitz # PyMuPDF
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parent.parent
|
||
|
|
INVOICES = ROOT / "invoices"
|
||
|
|
OCR_DIR = ROOT / "data" / "review" / "ocr"
|
||
|
|
|
||
|
|
# Find tesseract
|
||
|
|
TESSERACT = None
|
||
|
|
for p in ["C:/Program Files/Tesseract-OCR/tesseract.exe", "C:/Program Files (x86)/Tesseract-OCR/tesseract.exe"]:
|
||
|
|
if os.path.exists(p):
|
||
|
|
TESSERACT = p
|
||
|
|
break
|
||
|
|
if not TESSERACT:
|
||
|
|
sys.exit("tesseract.exe not found - install via winget install UB-Mannheim.TesseractOCR")
|
||
|
|
|
||
|
|
OCR_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
def ocr_page_image(img_bytes: bytes) -> str:
|
||
|
|
"""OCR a single PNG image, return extracted text."""
|
||
|
|
# Write to temp file
|
||
|
|
tmp_in = OCR_DIR / "_tmp_in.png"
|
||
|
|
tmp_out = OCR_DIR / "_tmp_out"
|
||
|
|
tmp_in.write_bytes(img_bytes)
|
||
|
|
try:
|
||
|
|
result = subprocess.run(
|
||
|
|
[TESSERACT, str(tmp_in), str(tmp_out), "-l", "eng"],
|
||
|
|
capture_output=True, text=True, timeout=60,
|
||
|
|
)
|
||
|
|
out_file = tmp_out.with_suffix(".txt")
|
||
|
|
if out_file.exists():
|
||
|
|
return out_file.read_text(encoding="utf-8", errors="replace")
|
||
|
|
return ""
|
||
|
|
finally:
|
||
|
|
tmp_in.unlink(missing_ok=True)
|
||
|
|
(tmp_out.with_suffix(".txt")).unlink(missing_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
def ocr_pdf(pdf_path: Path) -> list[str]:
|
||
|
|
"""Return list of OCR'd text per page."""
|
||
|
|
texts = []
|
||
|
|
with fitz.open(pdf_path) as doc:
|
||
|
|
for page_num, page in enumerate(doc, start=1):
|
||
|
|
# Render at 3x for better OCR accuracy
|
||
|
|
mat = fitz.Matrix(3, 3)
|
||
|
|
pix = page.get_pixmap(matrix=mat)
|
||
|
|
img_bytes = pix.tobytes("png")
|
||
|
|
text = ocr_page_image(img_bytes)
|
||
|
|
texts.append(text)
|
||
|
|
print(f" {pdf_path.name} page {page_num}/{len(doc)}: {len(text)} chars")
|
||
|
|
return texts
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
pdfs = sorted(p for p in INVOICES.rglob("*.pdf") if "_non_jei" not in p.parts)
|
||
|
|
errors = 0
|
||
|
|
for pdf in pdfs:
|
||
|
|
try:
|
||
|
|
out_dir = OCR_DIR / pdf.stem.strip()
|
||
|
|
out_dir.mkdir(exist_ok=True)
|
||
|
|
# OCR all pages
|
||
|
|
texts = ocr_pdf(pdf)
|
||
|
|
for i, text in enumerate(texts, start=1):
|
||
|
|
(out_dir / f"page_{i:03d}.txt").write_text(text, encoding="utf-8")
|
||
|
|
# Combined file
|
||
|
|
combined = "\n=== PAGE BREAK ===\n".join(texts)
|
||
|
|
(OCR_DIR / f"{pdf.stem.strip()}.txt").write_text(combined, encoding="utf-8")
|
||
|
|
except Exception as e:
|
||
|
|
errors += 1
|
||
|
|
print(f" ERR {pdf.name}: {e}")
|
||
|
|
print(f"\nDone. OCR text in {OCR_DIR}. Errors: {errors}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|