Update app.py

app.py

@@ -1,7 +1,15 @@
 """
-GLM-OCR Hugging Face Space app with client-side header/footer fallback for MaaS.
-Works for every PDF and every image: uses API bboxes when available, else minimal band.
+GLM-OCR Hugging Face Space app for PDF/image OCR with header inclusion
+and table-structure stabilization for downstream bank-statement pipelines.
+
+Hard-coded knobs (no environment variables required).
+
+Primary goal for reconcile rate:
+- preserve right-most columns (often "Balance") by higher DPI render + right padding
+- keep tables as tables (convert markdown pipe tables -> HTML table)
+- return ---page-separator--- between pages
 """
+
 # Patch asyncio first (before Gradio imports it) to suppress Python 3.13 cleanup noise
 import asyncio
 try:
@@ -18,29 +26,60 @@ except Exception:
 import logging
 import os
 import re
+import html
 import tempfile
+from typing import List, Tuple
+
 import yaml
 import gradio as gr
 
 import glmocr
 
 log = logging.getLogger("glmocr_app")
+logging.basicConfig(level=logging.INFO)
 
 GLMOCR_BASE = os.path.dirname(glmocr.__file__)
 CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml")
 FORMATTER_PATH = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 
-# When MaaS is enabled, the cloud API ignores local layout config and does not return
-# header/footer regions. We always run client-side OCR on top/bottom bands as fallback.
-DEFAULT_ZONE_FRAC = float(os.environ.get("GLMOCR_DEFAULT_ZONE_FRAC", "0.12"))
-# MaaS rejects very small images (400). Only send crops that meet minimum dimensions.
-MIN_CROP_HEIGHT = int(os.environ.get("GLMOCR_MIN_CROP_HEIGHT", "112"))
-MIN_CROP_PIXELS = int(os.environ.get("GLMOCR_MIN_CROP_PIXELS", "12544"))  # 112*112
-# Wider PDF bands for position-based extraction (capture account numbers etc. in top/bottom)
-PDF_HEADER_BAND_FRAC = float(os.environ.get("GLMOCR_PDF_HEADER_BAND", "0.08"))   # top 8%
-PDF_FOOTER_BAND_FRAC = float(os.environ.get("GLMOCR_PDF_FOOTER_BAND", "0.85"))   # bottom 15%
-
-# Single shared parser to avoid "GLM-OCR initialized" per request and asyncio cleanup issues.
+
+# ============================================================
+# HARD-CODED SETTINGS (edit these numbers to tune quality/speed)
+# ============================================================
+
+# 1) GLM-OCR MaaS API key
+# IMPORTANT: Do NOT hard-code secrets in a public Space.
+# If your Space is public, switch to HF Secrets instead.
+GLMOCR_API_KEY = "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w"
+
+# 2) Render quality (higher = better OCR for small/right-aligned digits; slower)
+RENDER_SCALE = 2.2   # try 2.5 if Balance column is still missing
+
+# 3) Add padding to protect columns near edges (Balance is usually right-most)
+PAD_LEFT_FRAC = 0.02
+PAD_RIGHT_FRAC = 0.06  # try 0.10 if right-most balances are missing
+PAD_TOP_FRAC = 0.01
+PAD_BOTTOM_FRAC = 0.01
+
+# 4) Mild contrast boost (helps faint gray text)
+ENABLE_CONTRAST = True
+
+# 5) Header/footer band heuristics
+DEFAULT_ZONE_FRAC = 0.12      # OCR band for header (top 12%) when regions not available
+PDF_HEADER_BAND_FRAC = 0.10   # PDF text fallback: take top 10% words if clip returns empty
+
+# Footer: disabled by default to avoid duplicating what GLM already returns
+ENABLE_FOOTER_OCR = False
+PDF_FOOTER_BAND_FRAC = 0.88   # bottom 12% (if footer enabled)
+
+# MaaS minimum image sizes for crops; we pad if needed
+MIN_CROP_HEIGHT = 112
+MIN_CROP_PIXELS = 112 * 112
+
+# ============================================================
+
+
+# Single shared parser to avoid re-init per request
 _parser = None
 
 def get_parser():
@@ -48,53 +87,26 @@ def get_parser():
     if _parser is None:
         from glmocr import GlmOcr
         _parser = GlmOcr(
-            api_key=os.environ.get("GLMOCR_API_KEY", "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w"),
+            api_key=GLMOCR_API_KEY,
             mode="maas",
         )
     return _parser
 
+
 # ---------------------------------------------------------------------------
-# 1. Config: set MaaS and optionally include headers/footers for non-MaaS
+# Best-effort config tweaks (safe to fail on read-only HF env)
 # ---------------------------------------------------------------------------
 try:
     with open(CONFIG_PATH, "r") as f:
         config = yaml.safe_load(f)
     config["pipeline"]["maas"]["enabled"] = True
-    config["pipeline"]["maas"]["api_key"] = os.environ.get("GLMOCR_API_KEY", "e2b138b2005a41cb9d87dd18805838aa.lyd51L23rcDbsw0w")
-
-    to_include_as_text = {"header", "footer"}
-    layout_section = config.get("pipeline", {}).get("layout", {})
-    if "label_task_mapping" in layout_section:
-        mapping = layout_section["label_task_mapping"]
-        abandon = mapping.get("abandon")
-        if isinstance(abandon, list):
-            mapping["abandon"] = [x for x in abandon if x not in to_include_as_text]
-        text_labels = mapping.get("text")
-        if isinstance(text_labels, list):
-            for label in to_include_as_text:
-                if label not in text_labels:
-                    text_labels.append(label)
-
-    formatter_section = config.get("pipeline", {}).get("result_formatter", {})
-    if "abandon" in formatter_section:
-        abandon = formatter_section["abandon"]
-        if isinstance(abandon, list):
-            formatter_section["abandon"] = [x for x in abandon if x not in to_include_as_text]
-    if "label_visualization_mapping" in formatter_section:
-        text_vis = formatter_section["label_visualization_mapping"].get("text")
-        if isinstance(text_vis, list):
-            for label in to_include_as_text:
-                if label not in text_vis:
-                    text_vis.append(label)
-
+    config["pipeline"]["maas"]["api_key"] = GLMOCR_API_KEY
     with open(CONFIG_PATH, "w") as f:
         yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 except Exception:
-    pass  # e.g. read-only package on Hugging Face; MaaS ignores layout anyway
+    pass
 
-# ---------------------------------------------------------------------------
-# 2. result_formatter: stop stripping header/footer (best-effort; may be read-only)
-# ---------------------------------------------------------------------------
+# Best-effort formatter tweak: avoid stripping header/footer labels
 try:
     with open(FORMATTER_PATH, "r") as f:
         source = f.read()
@@ -108,7 +120,11 @@ except Exception:
     pass
 
 
+# --------------------------
+# Header/footer helpers
+# --------------------------
 def get_header_footer_zones(regions, norm_height=1000):
+    """Infer header/footer extents from bbox regions if present."""
     if not regions:
         return None, None
     y_tops, y_bottoms = [], []
@@ -123,7 +139,7 @@ def get_header_footer_zones(regions, norm_height=1000):
 
 
 def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract text from a horizontal band using a clip rect."""
+    """Extract text from a horizontal band using a clip rect (works if PDF has text layer)."""
     try:
         import pymupdf as fitz
         doc = fitz.open(pdf_path)
@@ -138,8 +154,7 @@ def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
 
 
 def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
-    """Extract all text whose word bbox falls in the given vertical band (by position).
-    Use a wider band (e.g. top 15% / bottom 15%) to capture header/footer that clip might miss."""
+    """Extract words whose bbox intersects a vertical band (robust fallback)."""
     try:
         import pymupdf as fitz
         doc = fitz.open(pdf_path)
@@ -147,7 +162,7 @@ def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
         h = page.rect.height
         y_lo = h * y_start_frac
         y_hi = h * y_end_frac
-        words = page.get_text("words")  # list of (x0, y0, x1, y1, word, ...)
+        words = page.get_text("words")
         doc.close()
         parts = []
         for w in words:
@@ -161,7 +176,7 @@ def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
 
 
 def ocr_zone(image_path, y_start_frac, y_end_frac):
-    """Run OCR on a horizontal band. Pads small crops to meet API minimum size instead of skipping."""
+    """Run OCR on a horizontal band. Pads small crops to meet MaaS minimum size."""
     zone_name = "header" if y_end_frac < 0.5 else "footer"
     try:
         from PIL import Image
@@ -171,9 +186,10 @@ def ocr_zone(image_path, y_start_frac, y_end_frac):
         y1 = min(h, int(h * y_end_frac))
         if y1 <= y0:
             return ""
+
         crop = img.crop((0, y0, w, y1))
         cw, ch = crop.size
-        # Pad to minimum size so MaaS accepts the image (avoids 400 on tiny strips)
+
         if ch < MIN_CROP_HEIGHT or (cw * ch) < MIN_CROP_PIXELS:
             need_h = max(ch, MIN_CROP_HEIGHT)
             need_w = max(cw, 1)
@@ -181,11 +197,11 @@ def ocr_zone(image_path, y_start_frac, y_end_frac):
                 need_w = max(need_w, (MIN_CROP_PIXELS + need_h - 1) // need_h)
             canvas = Image.new("RGB", (need_w, need_h), (255, 255, 255))
             if zone_name == "header":
-                canvas.paste(crop, (0, 0))  # crop at top
+                canvas.paste(crop, (0, 0))
             else:
-                canvas.paste(crop, (0, need_h - ch))  # crop at bottom
+                canvas.paste(crop, (0, need_h - ch))
             crop = canvas
-            cw, ch = crop.size
+
         fd, path = tempfile.mkstemp(suffix=".jpg")
         os.close(fd)
         try:
@@ -195,8 +211,7 @@ def ocr_zone(image_path, y_start_frac, y_end_frac):
             if not isinstance(out, list):
                 out = [out]
             if out and getattr(out[0], "markdown_result", None):
-                text = (out[0].markdown_result or "").strip()
-                return text
+                return (out[0].markdown_result or "").strip()
         finally:
             try:
                 os.unlink(path)
@@ -207,48 +222,156 @@ def ocr_zone(image_path, y_start_frac, y_end_frac):
     return ""
 
 
-def fix_account_number(hdr):
-    """Fix header text:
-    1. If Account Number: has no value, find the number and inject it.
-    2. Remove any leading standalone number that is duplicated elsewhere in hdr.
-    """
-    import re
+def fix_account_number(hdr: str) -> str:
+    """Fix common account-number formatting issues."""
     if not hdr:
         return hdr
-    # Step 1: inject account number if missing after label
     if "Account Number:" in hdr and "Account Number: " not in hdr:
-        match = re.search("[0-9]{5,}", hdr)
-        if match:
-            acct = match.group(0)
-            hdr = hdr.replace("Account Number:", "Account Number: " + acct)
-    # Step 2: remove leading number that is duplicated after Account Number:
-    # e.g. "000000731823008 JPMorgan ... Account Number: 000000731823008"
-    # -> "JPMorgan ... Account Number: 000000731823008"
+        m = re.search(r"[0-9]{5,}", hdr)
+        if m:
+            hdr = hdr.replace("Account Number:", "Account Number: " + m.group(0))
     acct_match = re.search(r"Account Number: ([0-9]{5,})", hdr)
     if acct_match:
         acct = acct_match.group(1)
-        # If hdr starts with that same number, remove it from the start
         if hdr.startswith(acct):
             hdr = hdr[len(acct):].lstrip()
     return hdr
 
 
-def close_unclosed_html(md):
-    """Close any unclosed <table>, <tbody>, <tr> tags to prevent layout bleed."""
+# --------------------------
+# Table stabilization helpers
+# --------------------------
+def close_unclosed_html(md: str) -> str:
+    """Close unclosed <table>/<tr>/<td> tags to prevent bleed."""
     if not md:
         return md
-    import re
-    open_tags = re.findall(r'<(table|tbody|thead|tr|td|th)[\s>]', md)
-    close_tags = re.findall(r'</(table|tbody|thead|tr|td|th)>', md)
-    # Count unclosed tags and close them in reverse order
-    for tag in reversed(['table', 'tbody', 'thead', 'tr', 'td', 'th']):
-        opened = open_tags.count(tag)
-        closed = close_tags.count(tag)
+    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):
+        return sum(1 for t in tags if t.lower() == name)
+
+    for tag in reversed(["td", "th", "tr", "thead", "tbody", "table"]):
+        opened = count(open_tags, tag)
+        closed = count(close_tags, tag)
         if opened > closed:
-            md += '</{}>'.format(tag) * (opened - closed)
+            md += ("</%s>" % tag) * (opened - closed)
     return md
 
 
+def looks_like_markdown_table(block: str) -> bool:
+    """Detect simple markdown pipe tables."""
+    lines = [ln.rstrip() for ln in block.strip().splitlines() if ln.strip()]
+    if len(lines) < 2:
+        return False
+    if "|" not in lines[0]:
+        return False
+    sep = lines[1].replace(" ", "")
+    return ("---" in sep) and ("|" in sep)
+
+
+def md_table_to_html(block: str) -> str:
+    """Convert a simple markdown pipe table to HTML table (best-effort)."""
+    lines = [ln.strip() for ln in block.strip().splitlines() if ln.strip()]
+    if len(lines) < 2:
+        return block
+
+    def split_row(row: str):
+        row = row.strip()
+        if row.startswith("|"):
+            row = row[1:]
+        if row.endswith("|"):
+            row = row[:-1]
+        return [p.strip() for p in row.split("|")]
+
+    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>")
+    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>")
+    return "<table>\n" + "\n".join(html_rows) + "\n</table>"
+
+
+def normalize_money_glyphs(text: str) -> str:
+    """Conservative normalization for OCR number quirks."""
+    if not text:
+        return text
+    t = text.replace("−", "-").replace("–", "-").replace("—", "-")
+    t = re.sub(r"\(\s*\$?\s*([0-9]{1,3}(?:,[0-9]{3})*|[0-9]+)(\.[0-9]{2})\s*\)", r"-\1\2", t)
+
+    def o_to_zero(m):
+        token = m.group(0)
+        return token.replace("O", "0").replace("o", "0")
+    t = re.sub(r"\b[0-9Oo\$,.\-]{4,}\b", o_to_zero, t)
+    return t
+
+
+def stabilize_tables_and_text(page_md: str) -> str:
+    """Convert markdown pipe tables to HTML and close 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())
+    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)
+
+    stabilized = "\n\n".join(out_blocks)
+    return close_unclosed_html(stabilized)
+
+
+# --------------------------
+# PDF rendering with padding (critical for Balance column)
+# --------------------------
+def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
+    import pymupdf as fitz
+    from PIL import Image, ImageEnhance
+
+    doc = fitz.open(pdf_path)
+    page_images: List[str] = []
+    page_heights: List[int] = []
+
+    for i in range(len(doc)):
+        page = doc[i]
+        pix = page.get_pixmap(matrix=fitz.Matrix(RENDER_SCALE, RENDER_SCALE), alpha=False)
+
+        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+
+        if ENABLE_CONTRAST:
+            img = ImageEnhance.Contrast(img).enhance(1.12)
+
+        w, h = img.size
+        pad_l = int(w * PAD_LEFT_FRAC)
+        pad_r = int(w * PAD_RIGHT_FRAC)
+        pad_t = int(h * PAD_TOP_FRAC)
+        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.paste(img, (pad_l, pad_t))
+            img = canvas
+
+        img_path = os.path.join(tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{i}.png")
+        img.save(img_path, "PNG", compress_level=6)
+        page_images.append(img_path)
+        page_heights.append(img.height)
+
+    doc.close()
+    return page_images, page_heights
+
+
+# --------------------------
+# GLM-OCR result extraction
+# --------------------------
 def get_page_md_and_regions(page_result):
     md = ""
     if hasattr(page_result, "markdown_result") and page_result.markdown_result:
@@ -267,31 +390,27 @@ def get_page_md_and_regions(page_result):
     return md, regions
 
 
+# --------------------------
+# Main entry
+# --------------------------
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
-    try:
-        import pymupdf as fitz
 
+    page_images = []
+    try:
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
         is_pdf = path.lower().endswith(".pdf")
         parser = get_parser()
 
+        page_heights = []
+
         if is_pdf:
-            doc = fitz.open(path)
-            page_images = []
-            page_heights = []
-            for i in range(len(doc)):
-                pix = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
-                img_path = os.path.join(tempfile.gettempdir(), f"maas_page_{i}.png")
-                pix.save(img_path)
-                page_images.append(img_path)
-                page_heights.append(pix.height)
-            doc.close()
+            page_images, page_heights = render_pdf_pages_to_images(path)
             results = parser.parse(page_images)
         else:
             page_images = [path]
-            page_heights = []
+            page_heights = [1000]
             results = parser.parse(path)
 
         if not isinstance(results, list):
@@ -303,51 +422,68 @@ def run_ocr(uploaded_file):
             img_h = page_heights[page_num] if page_num < len(page_heights) else 1000
             header_end_frac, footer_start_frac = get_header_footer_zones(regions, img_h)
 
-            # MaaS does not return header/footer regions; always use at least default bands
             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)
-            if he <= 0:
-                he = DEFAULT_ZONE_FRAC  # ensure we always OCR a top band
-            if fs >= 1.0:
-                fs = 1.0 - DEFAULT_ZONE_FRAC  # ensure we always OCR a bottom band
+
+            # clamp
+            he = max(0.02, min(0.25, he))
+            fs = max(0.75, min(0.98, fs))
 
             parts = []
 
-            # Always run header band (top 8–12% of page)
-            if page_num < len(page_images):
-                img_path = page_images[page_num]
-                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)
+            # Header inclusion: PDF text -> band words -> OCR band
+            hdr = ""
+            if is_pdf:
+                hdr = extract_zone_text_pdf(path, page_num, 0, he)
                 if not (hdr and hdr.strip()):
-                    hdr = ocr_zone(img_path, 0, he)
-                if hdr and hdr.strip():
-                    parts.append(fix_account_number(hdr.strip()))
-
-            if page_md:
-                parts.append(close_unclosed_html(page_md))
-
-            # Footer extraction removed — MaaS API already captures all footer
-            # content correctly inside page_md. Extracting again causes duplication.
+                    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(fix_account_number(normalize_money_glyphs(hdr.strip())))
+
+            # Main OCR markdown, stabilized
+            if page_md and page_md.strip():
+                parts.append(stabilize_tables_and_text(page_md.strip()))
+
+            # Optional 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)
+                if not (ftr and ftr.strip()):
+                    ftr = ocr_zone(page_images[page_num], fs, 1.0)
+                if ftr and ftr.strip():
+                    parts.append(normalize_money_glyphs(ftr.strip()))
 
             if parts:
                 all_pages.append("\n\n".join(parts))
 
         return "\n\n---page-separator---\n\n".join(all_pages) if all_pages else "(No content)"
+
     except Exception as e:
         import traceback
         log.exception("run_ocr failed: %s", e)
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
+    finally:
+        # cleanup rendered images
+        for p in page_images:
+            try:
+                if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p):
+                    os.unlink(p)
+            except Exception:
+                pass
+
 
 with gr.Blocks(title="GLM-OCR") as demo:
-    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers and footers included.")
+    gr.Markdown("# GLM-OCR\nUpload a PDF or image. Headers included; tables stabilized.")
     file_in = gr.File(label="Upload PDF or image", file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"])
     run_btn = gr.Button("Run OCR", variant="primary")
-    out = gr.Textbox(lines=40, label="Output")
+    out = gr.Textbox(lines=40, label="Output (markdown)")
     run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()