Update app.py

app.py

@@ -4,6 +4,9 @@ import glob
 import io
 import os
 import re
+import shutil
+import subprocess
+import tempfile
 import unicodedata
 from collections import defaultdict
 
@@ -14,7 +17,7 @@ import pandas as pd
 import streamlit as st
 from docx import Document
 from langchain_groq import ChatGroq
-from PIL import Image, ImageFilter
+from PIL import Image, ImageOps
 from pptx import Presentation
 from rank_bm25 import BM25Okapi
 from sentence_transformers import SentenceTransformer
@@ -26,7 +29,9 @@ CHUNK_WORDS = 80          # stays under the embedding model's 128-token limit
 CHUNK_OVERLAP = 15
 TABLE_CHUNK_WORDS = 60    # table/spreadsheet rows are token-heavy, so smaller groups
 RRF_K = 60
-OCR_DPI = 400             # higher than before: small typewriter/fax text needs it. Lower (e.g. 300) for speed.
+OCR_DPI = 200             # Tesseract works best when text is ~30px tall; 400 hurt large poster text. Raise for tiny fax text.
+OCR_PSMS = (3, 11)        # 3 = auto page layout, 11 = sparse text (finds scattered words on posters)
+EXTRACT_VERSION = "ocr-cli-v3"  # bump to invalidate st.cache_data (it does not track helper functions)
 OCR_LANGS = "ara+eng"
 OCR_FALLBACK_LANGS = "eng"  # tried first (Arabic data can hurt English-only pages); also the fallback if Arabic data is missing
 OCR_CONTENT_THRESHOLD = 200  # gray level below which a pixel counts as "content" (0-255)
@@ -99,22 +104,26 @@ def _content_clip(page):
     except Exception:
         return page.rect
 
-def _prep_variants(page, clip):
-    """Render the page, flatten uneven background, stretch contrast, and also return an
-    inverted copy (for white-on-dark text such as banners). Returns [normal, inverted]."""
+def _render_gray(page, clip):
     pix = page.get_pixmap(dpi=OCR_DPI, clip=clip, colorspace=fitz.csGRAY)
-    img = Image.frombytes("L", (pix.width, pix.height), pix.samples)
-    bg = img.filter(ImageFilter.GaussianBlur(radius=max(15, OCR_DPI // 8)))
-    a = np.asarray(img, dtype=np.float32)
-    b = np.asarray(bg, dtype=np.float32) + 1.0
-    flat = np.clip(a / b * 255.0, 0, 255)  # divides out the uneven gray background
-    lo, hi = np.percentile(flat, (2, 98))
-    flat = np.clip((flat - lo) / max(hi - lo, 1) * 255.0, 0, 255).astype(np.uint8)
-    return [flat, 255 - flat]
-
-def _ocr_array(arr, lang):
-    h, w = arr.shape
-    pix = fitz.Pixmap(fitz.csGRAY, w, h, np.ascontiguousarray(arr).tobytes(), False)
+    return Image.frombytes("L", (pix.width, pix.height), pix.samples)
+
+def _tesseract_cli(img, lang, psm):
+    """Run the tesseract binary directly on a PIL image. Raises if it is missing or fails."""
+    with tempfile.TemporaryDirectory() as tmp:
+        path = os.path.join(tmp, "page.png")
+        img.save(path)
+        env = dict(os.environ)
+        r = subprocess.run(["tesseract", path, "stdout", "-l", lang, "--psm", str(psm)],
+                           capture_output=True, text=True, timeout=180, env=env)
+        if r.returncode != 0:
+            raise RuntimeError((r.stderr or "tesseract failed").strip()[:300])
+        return unicodedata.normalize("NFKC", r.stdout).strip()
+
+def _tesseract_pymupdf(img, lang):
+    """Fallback when the tesseract binary is not on PATH: PyMuPDF's built-in OCR."""
+    w, h = img.size
+    pix = fitz.Pixmap(fitz.csGRAY, w, h, img.tobytes(), False)
     pix.set_dpi(OCR_DPI, OCR_DPI)
     ocr_doc = fitz.open("pdf", pix.pdfocr_tobytes(language=lang))
     try:
@@ -124,39 +133,43 @@ def _ocr_array(arr, lang):
 
 def ocr_page(page):
     """OCR a scanned page. Returns (text, error).
-    Runs English-first and Arabic+English, each on a cleaned and an inverted image,
-    and merges the unique lines. error is None unless every attempt raised, so
-    'OCR ran and found nothing' can be told apart from 'OCR could not run'."""
+    Tries English first, then Arabic+English, each on the raw render and a contrast-stretched
+    inverted copy (for white-on-dark banners), with two page-segmentation modes, and merges the
+    unique lines. error is None unless every attempt raised."""
     _set_tessdata()
     clip = _content_clip(page)
     if clip is None:  # truly blank page: nothing to OCR, and not an error
-        return "", None
-    errors, lines, seen = [], [], set()
+        return "", None, "page looked blank (no pixels darker than the content threshold)"
     try:
-        variants = _prep_variants(page, clip)
+        raw = _render_gray(page, clip)
     except Exception as e:
-        return "", f"image preprocessing failed: {e}"
+        return "", f"could not render page: {e}", ""
+    inv = ImageOps.invert(ImageOps.autocontrast(raw, cutoff=2))
+    use_cli = shutil.which("tesseract") is not None
+    errors, lines, seen = [], [], set()
     for lang in dict.fromkeys((OCR_FALLBACK_LANGS, OCR_LANGS)):  # de-duplicated, order kept
-        for arr in variants:
-            try:
-                txt = _ocr_array(arr, lang)
-            except Exception as e:
-                errors.append(f"[{lang}] {e}")
-                continue
-            for ln in txt.splitlines():
-                key = re.sub(r"\W+", "", ln.lower())
-                if len(key) >= 3 and key not in seen:  # merge unique lines from all passes
-                    seen.add(key)
-                    lines.append(ln.strip())
+        for img in (raw, inv):
+            for psm in (OCR_PSMS if use_cli else (None,)):
+                try:
+                    txt = _tesseract_cli(img, lang, psm) if use_cli else _tesseract_pymupdf(img, lang)
+                except Exception as e:
+                    errors.append(f"[{lang}] {e}")
+                    continue
+                for ln in txt.splitlines():
+                    key = re.sub(r"\W+", "", ln.lower())
+                    if len(key) >= 3 and key not in seen:
+                        seen.add(key)
+                        lines.append(ln.strip())
+    info = f"engine={'tesseract CLI' if use_cli else 'PyMuPDF OCR'}, image={raw.size[0]}x{raw.size[1]}px, unique lines={len(lines)}"
     if lines:
-        return "\n".join(lines), None
-    return "", ("; ".join(errors) if errors else None)
+        return "\n".join(lines), None, info
+    return "", ("; ".join(errors) if errors else None), info
 
 def extract_pdf(name, data):
     blocks = []
     stats = {"pages": 0, "tables": 0, "notes": []}
     doc = fitz.open(stream=data, filetype="pdf")
-    empty_pages, ocr_pages, ocr_errors = [], [], []
+    empty_pages, ocr_pages, ocr_errors, ocr_infos = [], [], [], []
     for i, page in enumerate(doc, 1):
         stats["pages"] += 1
         loc = f"p.{i}"
@@ -165,7 +178,8 @@ def extract_pdf(name, data):
             blocks.append(make_block(tbl, name, loc, "table"))
             stats["tables"] += 1
         if len(text) < 50:
-            ocr_text, ocr_err = ocr_page(page)
+            ocr_text, ocr_err, ocr_info = ocr_page(page)
+            ocr_infos.append((i, len(ocr_text), ocr_info))
             if ocr_err:
                 ocr_errors.append((i, ocr_err))
             if len(ocr_text) > len(text):
@@ -186,6 +200,9 @@ def extract_pdf(name, data):
             "Check that Tesseract and its language data (eng, ara) are installed."
         )
     if empty_pages:
+        dbg = "; ".join(f"p.{p}: {n} chars, {inf}" for p, n, inf in ocr_infos if p in empty_pages[:3])
+        if dbg:
+            stats["notes"].append(f"OCR details: {dbg}")
         shown = ", ".join(map(str, empty_pages[:15])) + ("..." if len(empty_pages) > 15 else "")
         stats["notes"].append(
             f"Little or no readable text found on page(s) {shown}"
@@ -271,7 +288,7 @@ def extract_text_file(name, data):
     return blocks, {"pages": 1, "tables": 0, "notes": []}
 
 @st.cache_data(show_spinner=False)
-def extract_file(name, data):
+def extract_file(name, data, version=EXTRACT_VERSION):
     lower = name.lower()
     if lower.endswith(".pdf"):
         return extract_pdf(name, data)
@@ -386,6 +403,23 @@ with st.sidebar:
     min_rel = st.slider("Minimum relevance to answer", 0.0, 0.8, 0.25, 0.05,
                         help="If the best semantic match is below this, the app says 'not found' instead of guessing.")
 
+    with st.expander("OCR diagnostics"):
+        tbin = shutil.which("tesseract")
+        st.write(f"tesseract binary: `{tbin or 'NOT FOUND (using PyMuPDF OCR fallback)'}`")
+        if tbin:
+            try:
+                _set_tessdata()
+                out = subprocess.run(["tesseract", "--list-langs"], capture_output=True, text=True, timeout=20)
+                st.code((out.stdout or out.stderr).strip())
+            except Exception as e:
+                st.write(f"Could not query Tesseract: {e}")
+        st.write(f"TESSDATA_PREFIX: `{os.environ.get('TESSDATA_PREFIX', 'unset')}`")
+        st.caption(f"Extraction version: {EXTRACT_VERSION}")
+        if st.button("Clear extraction cache"):
+            st.cache_data.clear()
+            st.session_state.pop("signature", None)
+            st.rerun()
+
 embedder = load_embedder()
 
 groq_key = None
@@ -413,7 +447,7 @@ if uploaded_files:
         with st.spinner("Reading documents (scanned pages are OCR'd and may take a while)..."):
             for f in uploaded_files:
                 try:
-                    blocks, stats = extract_file(f.name, f.getvalue())
+                    blocks, stats = extract_file(f.name, f.getvalue(), EXTRACT_VERSION)
                 except Exception as e:
                     blocks, stats = [], {"pages": 0, "tables": 0, "notes": [f"Could not read file: {e}"]}
                 all_blocks.extend(blocks)