Add OCR fallback for scanned PDFs

src/codex_extractor.py

@@ -30,6 +30,12 @@ Use Prompt 2 (ChatGPT/Gemini) for those — see Codex Build Prompts document.
 Dependencies (add to requirements.txt):
   pdfplumber>=0.10
   nltk>=3.8
+  pypdfium2>=4.30
+  pytesseract>=0.3.10
+
+System dependency for OCR in Hugging Face Space (packages.txt):
+  tesseract-ocr
+  tesseract-ocr-eng
 
 NLTK data required (auto-downloaded on first run):
   punkt, punkt_tab, averaged_perceptron_tagger, cmudict, stopwords
@@ -46,6 +52,7 @@ import string
 import os
 from collections import Counter
 from pathlib import Path
+from shutil import which
 from typing import Any
 
 # ── OPTIONAL IMPORTS WITH GRACEFUL FALLBACK ──────────────────────────────────
@@ -56,6 +63,20 @@ try:
 except ImportError:
     PDF_AVAILABLE = False
 
+try:
+    import pypdfium2 as pdfium
+    PDFIUM_AVAILABLE = True
+except Exception:
+    PDFIUM_AVAILABLE = False
+
+try:
+    import pytesseract
+    from pytesseract import TesseractNotFoundError
+    PYTESSERACT_AVAILABLE = True
+except Exception:
+    PYTESSERACT_AVAILABLE = False
+    TesseractNotFoundError = RuntimeError  # type: ignore[assignment]
+
 try:
     import nltk
     # Auto-download required NLTK data if not present
@@ -109,17 +130,81 @@ QUESTION_PATTERN = re.compile(r"\?")
 
 # ── TEXT EXTRACTION ───────────────────────────────────────────────────────────
 
+def _word_count(text: str) -> int:
+    """Count alpha-ish tokens quickly for extraction health checks."""
+    return len(SIMPLE_TOKENISE_PATTERN.findall(text.lower()))
+
+
+def _ocr_runtime_ready() -> bool:
+    """Return True when OCR fallback dependencies and binary are available."""
+    if not PDFIUM_AVAILABLE or not PYTESSERACT_AVAILABLE:
+        return False
+    return which("tesseract") is not None
+
+
+def extract_text_from_pdf_ocr(pdf_path: str | Path) -> str:
+    """
+    OCR fallback for image-based PDFs.
+    Renders pages to images, then runs Tesseract OCR.
+    """
+    if not PDFIUM_AVAILABLE:
+        raise RuntimeError("OCR fallback unavailable: pypdfium2 is not installed.")
+    if not PYTESSERACT_AVAILABLE:
+        raise RuntimeError("OCR fallback unavailable: pytesseract is not installed.")
+    if which("tesseract") is None:
+        raise RuntimeError("OCR fallback unavailable: tesseract binary is not installed on this runtime.")
+
+    render_scale = float(os.getenv("OCR_RENDER_SCALE", "2.0"))
+    max_pages = int(os.getenv("OCR_MAX_PAGES", "200"))
+    ocr_lang = os.getenv("OCR_LANG", "eng")
+    ocr_config = os.getenv("OCR_CONFIG", "--oem 1 --psm 6")
+
+    text_parts: list[str] = []
+    pdf = pdfium.PdfDocument(str(pdf_path))
+    page_count = min(len(pdf), max_pages)
+
+    for page_index in range(page_count):
+        page = pdf[page_index]
+        bitmap = page.render(scale=render_scale)
+        pil_image = bitmap.to_pil()
+        page_text = pytesseract.image_to_string(pil_image, lang=ocr_lang, config=ocr_config)
+        if page_text and page_text.strip():
+            text_parts.append(page_text)
+
+    return "\n".join(text_parts)
+
+
 def extract_text_from_pdf(pdf_path: str | Path) -> str:
-    """Extract all text from a PDF file using pdfplumber."""
+    """
+    Extract text from a PDF:
+    1) native selectable text layer
+    2) OCR fallback if little/no selectable text is present
+    """
     if not PDF_AVAILABLE:
         raise RuntimeError("pdfplumber is not installed. Add it to requirements.txt.")
+
     text_parts = []
     with pdfplumber.open(str(pdf_path)) as pdf:
         for page in pdf.pages:
             page_text = page.extract_text()
             if page_text:
                 text_parts.append(page_text)
-    return "\n".join(text_parts)
+    extracted = "\n".join(text_parts)
+
+    # If selectable text exists, use it immediately.
+    if _word_count(extracted) >= 20:
+        return extracted
+
+    # Scanned/image-only PDFs fall back to OCR.
+    try:
+        ocr_text = extract_text_from_pdf_ocr(pdf_path)
+    except Exception as ocr_error:
+        raise RuntimeError(
+            "No selectable text detected in PDF and OCR fallback failed. "
+            f"Reason: {ocr_error}"
+        ) from ocr_error
+
+    return ocr_text
 
 
 def extract_text_from_file(file_path: str | Path) -> str:
@@ -691,7 +776,11 @@ def process_upload(
     try:
         raw_text = extract_text_from_file(file_path)
         if not raw_text or len(raw_text.split()) < 20:
-            return "ERROR: No usable text extracted from file. Check the PDF contains selectable text (not scanned images).", {}
+            return (
+                "ERROR: No usable text extracted from file. "
+                "If this is a scanned/image PDF, OCR could not recover enough text.",
+                {},
+            )
         fp = extract_fingerprint(
             text=raw_text,
             author_name=author_name,
@@ -700,6 +789,18 @@ def process_upload(
         )
         report = format_fingerprint_report(fp)
         return report, fp
+    except RuntimeError as e:
+        msg = str(e)
+        if "OCR fallback failed" in msg:
+            if not _ocr_runtime_ready():
+                return (
+                    "ERROR: This PDF appears image-based, but OCR runtime is not available yet. "
+                    "Install OCR dependencies (pytesseract, pypdfium2) and system package "
+                    "`tesseract-ocr` in the Space, then retry.",
+                    {},
+                )
+            return f"ERROR: OCR was attempted but failed: {msg}", {}
+        return f"ERROR: {msg}", {}
     except Exception as e:
         return f"ERROR: {type(e).__name__}: {str(e)}", {}