Update app.py

app.py

@@ -1,30 +1,3 @@
-#!/usr/bin/env python3
-"""
-Simplified GLM-OCR Hugging Face / local Gradio app.
-
-Scope (intentionally small):
-  - PDF → padded high-DPI page images → GLM-OCR body markdown
-  - Header band: PDF text extraction first, optional header OCR fallback
-  - Footer band: same pattern, with light dedup so we do not paste a full
-    transaction dump twice when the body already captured it
-
-Universal image pipeline (same for every PDF, no keywords / no bank logic):
-  - Higher rasterization scale + extra white padding so fine print, boxed
-    section labels, and right-aligned amounts sit farther from the clip edge.
-  - Mild contrast + unsharp mask on every raster sent to the model so
-    thin rules and small glyphs are easier to read before recognition.
-
-Explicitly omitted vs the heavy Space build:
-  - No text-layer row injection, institution-specific splits (UCB / Navy /
-    TD / First Horizon / …), doc-wide dedupe passes, or HTML table rewriting
-    keyed off column names / dates / amounts (layout fidelity comes from input
-    image quality, not post-hoc string rules).
-
-Configure GLMOCR_API_KEY (environment variable). Optional: glmocr + gradio +
-pymupdf + pillow installed.
-"""
-
-# Patch asyncio first (before Gradio imports it) to reduce Python 3.13 loop noise
 import asyncio
 
 try:
@@ -55,7 +28,7 @@ try:
     GLMOCR_BASE = os.path.dirname(glmocr.__file__)
     CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml")
 except ImportError:
-    glmocr = None  # type: ignore
+    glmocr = None
     GLMOCR_BASE = ""
     CONFIG_PATH = ""
 
@@ -63,57 +36,220 @@ log = logging.getLogger("glmocr_simple_app")
 logging.basicConfig(level=logging.INFO)
 
 # ---------------------------------------------------------------------------
-# Settings — tuned for dense financial PDFs; applies to every document
+# Settings
 # ---------------------------------------------------------------------------
 
 GLMOCR_API_KEY = "cee1d52dd91a4ab591b3f6e105f8ad89.LgbQTECuzX0zrito"
 if not GLMOCR_API_KEY:
-    log.warning("GLMOCR_API_KEY is not set; GlmOcr() will fail until you export it.")
+    log.warning("GLMOCR_API_KEY is not set.")
 
-# Rasterization: higher scale = more pixels per PDF point (helps small type,
-# boxed headers, and narrow columns). Same constant for all uploads.
 RENDER_SCALE = 3.05
 
-# White margin as a fraction of page width/height after render. Extra right
-# margin helps right-aligned currency columns that hug the page edge.
-PAD_LEFT_FRAC = 0.035
-PAD_RIGHT_FRAC = 0.11
-PAD_TOP_FRAC = 0.018
+PAD_LEFT_FRAC   = 0.035
+PAD_RIGHT_FRAC  = 0.11
+PAD_TOP_FRAC    = 0.018
 PAD_BOTTOM_FRAC = 0.018
 
-ENABLE_CONTRAST = True
-# Slight contrast lift only; same factor for every file.
-CONTRAST_FACTOR = 1.16
+ENABLE_CONTRAST  = True
+CONTRAST_FACTOR  = 1.16
 
-# Subtle edge enhancement after contrast (helps hairlines and small digits).
-ENABLE_UNSHARP = True
-UNSHARP_RADIUS = 0.78
-UNSHARP_PERCENT = 72
+ENABLE_UNSHARP    = True
+UNSHARP_RADIUS    = 0.78
+UNSHARP_PERCENT   = 72
 UNSHARP_THRESHOLD = 1
 
-DEFAULT_ZONE_FRAC = 0.12
+DEFAULT_ZONE_FRAC    = 0.12
 PDF_HEADER_BAND_FRAC = 0.10
 
-ENABLE_FOOTER_OCR = True
+ENABLE_FOOTER_OCR    = True
 PDF_FOOTER_BAND_FRAC = 0.88
 
 MIN_CROP_HEIGHT = 112
 MIN_CROP_PIXELS = 112 * 112
 
-# PNG compression 0–9; lower = less loss before GLM-OCR (same for all PDFs).
 PAGE_PNG_COMPRESS_LEVEL = 3
-# JPEG quality for small header/footer crops sent to the API.
-ZONE_JPEG_QUALITY = 95
+ZONE_JPEG_QUALITY       = 95
+
+# ---------------------------------------------------------------------------
+# Similarity threshold for header deduplication (0-1). Texts with a
+# normalised token-overlap above this value are considered repeating headers
+# and suppressed on pages > 0.  No document keywords are used — works purely
+# on character/token similarity so it generalises to any PDF layout.
+# ---------------------------------------------------------------------------
+HEADER_SIMILARITY_THRESHOLD = 0.60
+
+# Minimum ratio of shared tokens (vs shorter string) to treat two table
+# schemas as "the same" and attempt column-count normalisation.
+TABLE_SCHEMA_SIMILARITY = 0.70
 
 _parser = None
 
 
-def _enhance_raster_for_ocr(img):
+# ---------------------------------------------------------------------------
+# Generic text-similarity helper (no hardcoding, no document keywords)
+# ---------------------------------------------------------------------------
+
+def _token_similarity(a: str, b: str) -> float:
+    """
+    Return a 0-1 Jaccard-style similarity between two strings based on
+    whitespace-split token sets.  Used to detect repeated headers and
+    identical table schema rows across pages without any hardcoded patterns.
+    """
+    if not a or not b:
+        return 0.0
+    ta = set(a.lower().split())
+    tb = set(b.lower().split())
+    if not ta or not tb:
+        return 0.0
+    return len(ta & tb) / len(ta | tb)
+
+
+def _normalise_text(t: str) -> str:
+    """Collapse whitespace, strip punctuation edges — for similarity checks."""
+    return re.sub(r"\s+", " ", t).strip()
+
+
+# ---------------------------------------------------------------------------
+# Generic HTML table normalisation
+# ---------------------------------------------------------------------------
+
+def _parse_html_table_rows(table_html: str):
+    """
+    Parse any HTML table into a list of rows, where each row is a list of
+    (text, colspan) tuples.  Works regardless of whether the source uses
+    <thead>/<tbody>, inline border attributes, or Bootstrap classes.
+    No document-specific assumptions.
+    """
+    rows = []
+    for row_m in re.finditer(r"<tr\b[^>]*>(.*?)</tr>", table_html, re.IGNORECASE | re.DOTALL):
+        cells = []
+        for cell_m in re.finditer(
+            r"<t[dh]\b([^>]*)>(.*?)</t[dh]>", row_m.group(1), re.IGNORECASE | re.DOTALL
+        ):
+            attrs, content = cell_m.group(1), cell_m.group(2)
+            cs_m = re.search(r'colspan\s*=\s*["\']?(\d+)', attrs, re.IGNORECASE)
+            colspan = int(cs_m.group(1)) if cs_m else 1
+            text = re.sub(r"<[^>]+>", " ", content)
+            text = re.sub(r"\s+", " ", text).strip()
+            cells.append((text, colspan))
+        if cells:
+            rows.append(cells)
+    return rows
+
+
+def _effective_col_count(rows) -> int:
+    """Return the most common effective column count across all rows."""
+    from collections import Counter
+    counts = Counter(sum(cs for _, cs in r) for r in rows)
+    return counts.most_common(1)[0][0] if counts else 0
+
+
+def _rows_to_uniform_html(rows, col_count: int) -> str:
+    """
+    Re-serialise parsed rows back to a clean, uniform HTML table where every
+    data row has exactly col_count <td> cells.  Header rows (those that had
+    all-<th> content in the original) are written as <th>.
+    Collapsed cells are split back to individual cells with empty padding so
+    downstream parsers always see the same schema.
+    """
+    out = ["<table>"]
+    for i, row in enumerate(rows):
+        effective = sum(cs for _, cs in row)
+        tag = "th" if i == 0 else "td"
+        out.append("<tr>")
+        filled = 0
+        for text, cs in row:
+            out.append(f"<{tag}>{html.escape(text)}</{tag}>")
+            filled += cs
+        # Pad missing columns
+        while filled < col_count:
+            out.append(f"<{tag}></{tag}>")
+            filled += 1
+        out.append("</tr>")
+    out.append("</table>")
+    return "\n".join(out)
+
+
+def normalise_tables(md: str) -> str:
+    """
+    Find all HTML tables in a markdown string and rewrite them so that:
+      1. Every table uses a uniform schema (same col count in every row).
+      2. Mixed <thead>/<tbody>/border="1"/class="..." attributes are stripped
+         to a clean, consistent <table> with plain <tr><th>/<td> cells.
+      3. colspan cells that inflate or deflate the logical column count are
+         expanded back to individual cells.
+
+    This is fully generic — no document keywords, no hardcoded column counts.
+    """
+    def replace_table(m):
+        raw = m.group(0)
+        rows = _parse_html_table_rows(raw)
+        if not rows:
+            return raw
+        col_count = _effective_col_count(rows)
+        if col_count == 0:
+            return raw
+        return _rows_to_uniform_html(rows, col_count)
+
+    # Match any HTML table, including those with class/border attrs and
+    # nested thead/tbody, non-greedily.
+    return re.sub(
+        r"<table\b[^>]*>.*?</table>",
+        replace_table,
+        md,
+        flags=re.IGNORECASE | re.DOTALL,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Image placeholder removal
+# ---------------------------------------------------------------------------
+
+def _strip_image_placeholders(md: str) -> str:
     """
-    Improve legibility of every raster passed to GLM-OCR (full pages and
-    header/footer crops). No document text or keywords — same pipeline for
-    all PDFs and images.
+    Remove OCR pipeline artefact image tags like ![Image N-N](imgs/...).
+    These reference local temp paths that never exist in the output context.
+    Generic regex — works for any img reference with a local relative path.
     """
+    return re.sub(r"!\[[^\]]*\]\(imgs/[^)]+\)", "", md)
+
+
+# ---------------------------------------------------------------------------
+# Generic repeated-header suppression
+# ---------------------------------------------------------------------------
+
+class _HeaderTracker:
+    """
+    Tracks header texts seen across pages using token-similarity comparison.
+    On the first occurrence a header is stored and emitted.  On subsequent
+    pages, if the candidate header is sufficiently similar to ANY previously
+    seen header it is suppressed entirely.
+
+    No document keywords.  Works for any PDF where the same block of text
+    repeats at the top of every page (date ranges, account numbers, titles,
+    report headers, etc.).
+    """
+
+    def __init__(self, threshold: float = HEADER_SIMILARITY_THRESHOLD):
+        self._seen: List[str] = []
+        self._threshold = threshold
+
+    def should_include(self, hdr: str) -> bool:
+        norm = _normalise_text(hdr)
+        if not norm:
+            return False
+        for seen in self._seen:
+            if _token_similarity(norm, seen) >= self._threshold:
+                return False
+        self._seen.append(norm)
+        return True
+
+
+# ---------------------------------------------------------------------------
+# Image enhancement
+# ---------------------------------------------------------------------------
+
+def _enhance_raster_for_ocr(img):
     from PIL import ImageEnhance, ImageFilter
 
     if ENABLE_CONTRAST:
@@ -135,7 +271,6 @@ def get_parser():
         raise RuntimeError("glmocr is not installed.")
     if _parser is None:
         from glmocr import GlmOcr
-
         _parser = GlmOcr(api_key=GLMOCR_API_KEY, mode="maas")
     return _parser
 
@@ -262,14 +397,14 @@ def fix_account_number(hdr: str) -> str:
     if acct_match:
         acct = acct_match.group(1)
         if hdr.startswith(acct):
-            hdr = hdr[len(acct) :].lstrip()
+            hdr = hdr[len(acct):].lstrip()
     return hdr
 
 
 def close_unclosed_html(md: str) -> str:
     if not md:
         return md
-    open_tags = re.findall(r"<(table|tbody|thead|tr|td|th)\b", md, flags=re.IGNORECASE)
+    open_tags  = re.findall(r"<(table|tbody|thead|tr|td|th)\b",  md, flags=re.IGNORECASE)
     close_tags = re.findall(r"</(table|tbody|thead|tr|td|th)>", md, flags=re.IGNORECASE)
 
     def count(tags, name):
@@ -306,17 +441,21 @@ def md_table_to_html(block: str) -> str:
             row = row[:-1]
         return [p.strip() for p in row.split("|")]
 
-    header = split_row(lines[0])
+    header     = split_row(lines[0])
     body_lines = [ln for ln in lines[2:] if "|" in ln]
 
     html_rows = []
-    html_rows.append("<tr>" + "".join(f"<th>{html.escape(c)}</th>" for c in header) + "</tr>")
+    html_rows.append(
+        "<tr>" + "".join(f"<th>{html.escape(c)}</th>" for c in header) + "</tr>"
+    )
     for ln in body_lines:
         cols = split_row(ln)
         if len(cols) < len(header):
             cols += [""] * (len(header) - len(cols))
         html_rows.append(
-            "<tr>" + "".join(f"<td>{html.escape(c)}</td>" for c in cols[: len(header)]) + "</tr>"
+            "<tr>"
+            + "".join(f"<td>{html.escape(c)}</td>" for c in cols[: len(header)])
+            + "</tr>"
         )
     return "<table>\n" + "\n".join(html_rows) + "\n</table>"
 
@@ -340,18 +479,36 @@ def normalize_money_glyphs(text: str) -> str:
 
 
 def light_stabilize_markdown(page_md: str) -> str:
-    """Convert obvious GitHub-style pipe tables to HTML; normalize money glyphs; repair tags."""
+    """
+    1. Convert pipe tables to HTML.
+    2. Normalise money glyphs.
+    3. Normalise all HTML table schemas to uniform column counts.
+    4. Strip broken image placeholder tags.
+    5. Repair unclosed HTML tags.
+    """
     if not page_md:
         return page_md
+
     page_md = normalize_money_glyphs(page_md)
-    blocks = re.split(r"\n\s*\n", page_md.strip())
+
+    # Convert pipe-style markdown tables to HTML first
+    blocks    = re.split(r"\n\s*\n", page_md.strip())
     out_blocks = []
     for b in blocks:
         if looks_like_markdown_table(b):
             out_blocks.append(md_table_to_html(b))
         else:
             out_blocks.append(b)
-    return close_unclosed_html("\n\n".join(out_blocks))
+    page_md = "\n\n".join(out_blocks)
+
+    # Normalise all HTML table schemas (fixes mixed colspan, thead/tbody,
+    # border="1" vs class="table", 3-col vs 4-col inconsistencies)
+    page_md = normalise_tables(page_md)
+
+    # Remove broken local image references left by the OCR region detector
+    page_md = _strip_image_placeholders(page_md)
+
+    return close_unclosed_html(page_md)
 
 
 def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
@@ -364,7 +521,9 @@ def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
 
     for i in range(len(doc)):
         page = doc[i]
-        pix = page.get_pixmap(matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False)
+        pix = page.get_pixmap(
+            matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False
+        )
 
         img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
         img = _enhance_raster_for_ocr(img)
@@ -376,11 +535,15 @@ def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
         pad_b = int(h * PAD_BOTTOM_FRAC)
 
         if any(p > 0 for p in (pad_l, pad_r, pad_t, pad_b)):
-            canvas = Image.new("RGB", (w + pad_l + pad_r, h + pad_t + pad_b), (255, 255, 255))
+            canvas = Image.new(
+                "RGB", (w + pad_l + pad_r, h + pad_t + pad_b), (255, 255, 255)
+            )
             canvas.paste(img, (pad_l, pad_t))
             img = canvas
 
-        img_path = os.path.join(tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{i}.png")
+        img_path = os.path.join(
+            tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{i}.png"
+        )
         img.save(img_path, "PNG", compress_level=PAGE_PNG_COMPRESS_LEVEL)
         page_images.append(img_path)
         page_heights.append(img.height)
@@ -413,9 +576,9 @@ def run_ocr(uploaded_file):
 
     page_images: List[str] = []
     try:
-        path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
-        is_pdf = path.lower().endswith(".pdf")
-        parser = get_parser()
+        path    = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
+        is_pdf  = path.lower().endswith(".pdf")
+        parser  = get_parser()
 
         page_heights: List[int] = []
 
@@ -423,13 +586,19 @@ def run_ocr(uploaded_file):
             page_images, page_heights = render_pdf_pages_to_images(path)
             results = parser.parse(page_images)
         else:
-            page_images = [path]
+            page_images  = [path]
             page_heights = [1000]
-            results = parser.parse(path)
+            results      = parser.parse(path)
 
         if not isinstance(results, list):
             results = [results]
 
+        # ------------------------------------------------------------------
+        # One tracker per document run — suppresses repeated headers across
+        # ALL pages without any hardcoded pattern matching.
+        # ------------------------------------------------------------------
+        header_tracker = _HeaderTracker(threshold=HEADER_SIMILARITY_THRESHOLD)
+
         all_pages = []
         for page_num, page_result in enumerate(results):
             page_md, regions = get_page_md_and_regions(page_result)
@@ -437,39 +606,68 @@ def run_ocr(uploaded_file):
             header_end_frac, footer_start_frac = get_header_footer_zones(regions, img_h)
 
             he = header_end_frac if header_end_frac is not None else DEFAULT_ZONE_FRAC
-            fs = footer_start_frac if footer_start_frac is not None else (1.0 - DEFAULT_ZONE_FRAC)
+            fs = (
+                footer_start_frac
+                if footer_start_frac is not None
+                else (1.0 - DEFAULT_ZONE_FRAC)
+            )
 
             he = max(0.02, min(0.25, he))
             fs = max(0.75, min(0.98, fs))
 
             parts = []
 
+            # --------------------------------------------------------------
+            # Header extraction with cross-page deduplication.
+            # The tracker uses token-similarity — no document-specific logic.
+            # --------------------------------------------------------------
             hdr = ""
             if is_pdf:
                 hdr = extract_zone_text_pdf(path, page_num, 0, he)
                 if not (hdr and hdr.strip()):
-                    hdr = extract_pdf_text_in_band(path, page_num, 0, PDF_HEADER_BAND_FRAC)
+                    hdr = extract_pdf_text_in_band(
+                        path, page_num, 0, PDF_HEADER_BAND_FRAC
+                    )
             if not (hdr and hdr.strip()) and page_num < len(page_images):
                 hdr = ocr_zone(page_images[page_num], 0, he)
+
             if hdr and hdr.strip():
-                parts.append(light_stabilize_markdown(fix_account_number(normalize_money_glyphs(hdr.strip()))))
+                # Only include this header if it is meaningfully different
+                # from any header already seen in this document.
+                if header_tracker.should_include(hdr.strip()):
+                    cleaned_hdr = light_stabilize_markdown(
+                        fix_account_number(
+                            normalize_money_glyphs(hdr.strip())
+                        )
+                    )
+                    if cleaned_hdr:
+                        parts.append(cleaned_hdr)
 
+            # Body
             if page_md and page_md.strip():
                 parts.append(light_stabilize_markdown(page_md.strip()))
 
+            # Footer
             if ENABLE_FOOTER_OCR and page_num < len(page_images):
                 ftr = ""
                 if is_pdf:
                     ftr = extract_zone_text_pdf(path, page_num, fs, 1.0)
                     if not (ftr and ftr.strip()):
-                        ftr = extract_pdf_text_in_band(path, page_num, PDF_FOOTER_BAND_FRAC, 1.0)
+                        ftr = extract_pdf_text_in_band(
+                            path, page_num, PDF_FOOTER_BAND_FRAC, 1.0
+                        )
                 if not (ftr and ftr.strip()):
                     ftr = ocr_zone(page_images[page_num], fs, 1.0)
+
                 if ftr and ftr.strip():
                     ftr_clean = normalize_money_glyphs(ftr.strip())
 
                     ftr_first_line = next(
-                        (ln.strip().lower() for ln in ftr_clean.splitlines() if ln.strip()),
+                        (
+                            ln.strip().lower()
+                            for ln in ftr_clean.splitlines()
+                            if ln.strip()
+                        ),
                         "",
                     )
                     already_present = ftr_first_line and any(
@@ -477,9 +675,9 @@ def run_ocr(uploaded_file):
                     )
 
                     _footer_date_re = re.compile(r"\b\d{1,2}[-/]\d{2}\b")
-                    _footer_amt_re = re.compile(r"\b\d{1,3}(?:,\d{3})*\.\d{2}\b")
+                    _footer_amt_re  = re.compile(r"\b\d{1,3}(?:,\d{3})*\.\d{2}\b")
                     _date_hits = len(_footer_date_re.findall(ftr_clean))
-                    _amt_hits = len(_footer_amt_re.findall(ftr_clean))
+                    _amt_hits  = len(_footer_amt_re.findall(ftr_clean))
                     is_txn_dump = _date_hits >= 3 and _amt_hits >= 3
 
                     if not already_present and not is_txn_dump:
@@ -488,7 +686,11 @@ def run_ocr(uploaded_file):
             if parts:
                 all_pages.append("\n\n".join(parts))
 
-        merged = "\n\n---page-separator---\n\n".join(all_pages) if all_pages else "(No content)"
+        merged = (
+            "\n\n---page-separator---\n\n".join(all_pages)
+            if all_pages
+            else "(No content)"
+        )
         return merged
 
     except Exception as e:
@@ -500,7 +702,11 @@ def run_ocr(uploaded_file):
     finally:
         for p in page_images:
             try:
-                if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p):
+                if (
+                    isinstance(p, str)
+                    and p.endswith(".png")
+                    and "glmocr_page_" in os.path.basename(p)
+                ):
                     os.unlink(p)
             except Exception:
                 pass
@@ -520,10 +726,10 @@ def _create_gradio_demo():
             file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
         )
         run_btn = gr.Button("Run OCR", variant="primary")
-        out = gr.Textbox(lines=40, label="Output (markdown / light HTML)")
+        out     = gr.Textbox(lines=40, label="Output (markdown / light HTML)")
         run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
     return demo
 
 
 if __name__ == "__main__":
-    _create_gradio_demo().launch()
+    _create_gradio_demo().launch()