Restore latest extractor with OCR fallback

src/codex_extractor.py

@@ -69,6 +69,7 @@ import re
 import string
 from collections import Counter
 from pathlib import Path
+from shutil import which
 from typing import Any
 
 # ── OPTIONAL IMPORTS WITH GRACEFUL FALLBACK ──────────────────────────────────
@@ -79,6 +80,18 @@ try:
 except ImportError:
     PDF_AVAILABLE = False
 
+try:
+    import pypdfium2 as pdfium
+    PDFIUM_AVAILABLE = True
+except Exception:
+    PDFIUM_AVAILABLE = False
+
+try:
+    import pytesseract
+    PYTESSERACT_AVAILABLE = True
+except Exception:
+    PYTESSERACT_AVAILABLE = False
+
 try:
     import nltk
     _NLTK_DATA = ["punkt", "punkt_tab", "averaged_perceptron_tagger", "cmudict", "stopwords"]
@@ -221,6 +234,67 @@ def is_artefact_line(line: str) -> bool:
 
 # ── TEXT EXTRACTION ───────────────────────────────────────────────────────────
 
+def _word_count(text: str) -> int:
+    return len(SIMPLE_TOKENISE_PATTERN.findall((text or "").lower()))
+
+
+def _ocr_runtime_ready() -> bool:
+    return PDFIUM_AVAILABLE and PYTESSERACT_AVAILABLE and which("tesseract") is not None
+
+
+def extract_text_from_pdf_ocr(
+    pdf_path: str | Path,
+    start_page: int | None = None,
+    end_page: int | None = None,
+) -> tuple[str, str]:
+    """
+    OCR fallback for scanned/image-only PDFs.
+    Returns (story_text, raw_text) using the same page-filter logic as native extraction.
+    """
+    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.")
+
+    render_scale = float(os.getenv("OCR_RENDER_SCALE", "2.0"))
+    ocr_lang = os.getenv("OCR_LANG", "eng")
+    ocr_config = os.getenv("OCR_CONFIG", "--oem 1 --psm 6")
+    max_pages = int(os.getenv("OCR_MAX_PAGES", "300"))
+
+    raw_parts: list[str] = []
+    story_parts: list[str] = []
+
+    pdf = pdfium.PdfDocument(str(pdf_path))
+    total_pages = min(len(pdf), max_pages)
+
+    idx_start = (start_page - 1) if start_page is not None else 0
+    idx_end = (end_page - 1) if end_page is not None else (total_pages - 1)
+    idx_start = max(0, idx_start)
+    idx_end = min(total_pages - 1, idx_end)
+
+    for page_idx in range(total_pages):
+        page = pdf[page_idx]
+        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) or ""
+        raw_parts.append(page_text)
+
+        if start_page is not None or end_page is not None:
+            if idx_start <= page_idx <= idx_end:
+                story_parts.append(page_text)
+            continue
+
+        if total_pages > 1 and page_idx == 0:
+            continue
+        if is_front_matter_page(page_text):
+            continue
+        story_parts.append(page_text)
+
+    return "\n".join(story_parts), "\n".join(raw_parts)
+
+
 def extract_text_from_pdf(
     pdf_path: str | Path,
     start_page: int | None = None,
@@ -270,7 +344,7 @@ def extract_text_from_pdf(
 
             # Automatic front-matter detection (Fix 2).
             # Always skip the first page (cover).
-            if page_idx == 0:
+            if total_pages > 1 and page_idx == 0:
                 continue
             if is_front_matter_page(page_text):
                 continue
@@ -279,6 +353,24 @@ def extract_text_from_pdf(
 
     raw_text = "\n".join(raw_parts)
     story_text = "\n".join(story_parts)
+
+    # If native text-layer extraction looks healthy, use it.
+    if _word_count(story_text) >= 20:
+        return story_text, raw_text
+
+    # Otherwise, try OCR fallback for scanned/image-only pages.
+    try:
+        ocr_story_text, ocr_raw_text = extract_text_from_pdf_ocr(
+            pdf_path,
+            start_page=start_page,
+            end_page=end_page,
+        )
+        if _word_count(ocr_story_text) >= max(20, _word_count(story_text)):
+            return ocr_story_text, ocr_raw_text
+    except Exception:
+        # Fall through: process_upload will raise a clear error if text is still insufficient.
+        pass
+
     return story_text, raw_text
 
 
@@ -1303,10 +1395,17 @@ def process_upload(
         )
 
         if not story_text or len(story_text.split()) < 20:
+            if str(file_path).lower().endswith(".pdf") and not _ocr_runtime_ready():
+                return (
+                    "ERROR: No usable text extracted from file and OCR runtime is unavailable. "
+                    "Install OCR dependencies (`pypdfium2`, `pytesseract`) and system package "
+                    "`tesseract-ocr` in the Space build.",
+                    {},
+                )
             return (
                 "ERROR: No usable text extracted from file. "
-                "Check the PDF contains selectable text (not scanned images), "
-                "or try specifying a story page range.",
+                "OCR fallback could not recover enough text. "
+                "Try a cleaner scan, higher resolution pages, or a story page range.",
                 {},
             )