Update app.py

app.py

@@ -1,46 +1,17 @@
 #!/usr/bin/env python3
 
-import asyncio
 import logging
 import os
-import re
 import tempfile
 import uuid
 from typing import List, Tuple
 
-import yaml
-
-try:
-    _orig_close = asyncio.BaseEventLoop.close
-
-    def _safe_close(self):
-        try:
-            _orig_close(self)
-        except (ValueError, OSError):
-            pass
-
-    asyncio.BaseEventLoop.close = _safe_close
-except Exception:
-    pass
-
-try:
-    import glmocr
-
-    GLMOCR_BASE = os.path.dirname(glmocr.__file__)
-    CONFIG_PATH = os.path.join(GLMOCR_BASE, "config.yaml")
-except ImportError:
-    glmocr = None
-    GLMOCR_BASE = ""
-    CONFIG_PATH = ""
-
 log = logging.getLogger("glmocr_simple_app")
 logging.basicConfig(level=logging.INFO)
 
-GLMOCR_API_KEY = "cee1d52dd91a4ab591b3f6e105f8ad89.LgbQTECuzX0zrito"
-if not GLMOCR_API_KEY:
-    log.warning(
-        "No ZHIPU_API_KEY or GLMOCR_API_KEY in environment; GlmOcr() will fail until you set one."
-    )
+# ── Fine-tuned model repo on HuggingFace ─────────────────────────────────────
+# Loads from HF Hub at runtime — no local storage needed in the Space
+MERGED_MODEL_DIR = os.environ.get("MODEL_DIR", "SimpleCodeAI/glm-ocr-finetuned")
 
 RENDER_SCALE = 3.0
 PAD_LEFT_FRAC = 0.035
@@ -55,17 +26,36 @@ UNSHARP_RADIUS = 0.78
 UNSHARP_PERCENT = 76
 UNSHARP_THRESHOLD = 1
 
-DEFAULT_ZONE_FRAC = 0.12
-PDF_HEADER_BAND_FRAC = 0.10
-ENABLE_FOOTER_OCR = True
-PDF_FOOTER_BAND_FRAC = 0.88
-
-MIN_CROP_HEIGHT = 112
-MIN_CROP_PIXELS = 112 * 112
 PAGE_PNG_COMPRESS_LEVEL = 3
-ZONE_JPEG_QUALITY = 95
+MAX_NEW_TOKENS = 2048
+MAX_IMAGE_SIDE = 1344  # resize longest side to this before inference
 
-_parser = None
+# ── Model singleton ───────────────────────────────────────────────────────────
+_model = None
+_processor = None
+
+
+def _load_model():
+    global _model, _processor
+    if _model is not None:
+        return _model, _processor
+
+    import torch
+    from transformers import AutoProcessor, AutoModelForImageTextToText
+
+    log.info("Loading fine-tuned model from %s ...", MERGED_MODEL_DIR)
+    _processor = AutoProcessor.from_pretrained(
+        MERGED_MODEL_DIR, trust_remote_code=True
+    )
+    _model = AutoModelForImageTextToText.from_pretrained(
+        MERGED_MODEL_DIR,
+        dtype=torch.bfloat16,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    _model.eval()
+    log.info("Model loaded.")
+    return _model, _processor
 
 
 def _enhance_raster_for_ocr(img):
@@ -84,129 +74,67 @@ def _enhance_raster_for_ocr(img):
     return img
 
 
-def get_parser():
-    global _parser
-    if glmocr is None:
-        raise RuntimeError("glmocr is not installed.")
-    if _parser is None:
-        from glmocr import GlmOcr
+def _resize_for_inference(img):
+    """Resize image preserving aspect ratio so longest side <= MAX_IMAGE_SIDE."""
+    from PIL import Image
 
-        kw = {"mode": "maas"}
-        if GLMOCR_API_KEY:
-            kw["api_key"] = GLMOCR_API_KEY
-        _parser = GlmOcr(**kw)
-    return _parser
+    w, h = img.size
+    longest = max(w, h)
+    if longest <= MAX_IMAGE_SIDE:
+        return img
+    ratio = MAX_IMAGE_SIDE / longest
+    new_size = (int(w * ratio), int(h * ratio))
+    return img.resize(new_size, Image.LANCZOS)
 
 
-if CONFIG_PATH:
-    try:
-        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
-            config = yaml.safe_load(f)
-        config.setdefault("pipeline", {}).setdefault("maas", {})
-        config["pipeline"]["maas"]["enabled"] = True
-        config["pipeline"]["maas"]["api_key"] = GLMOCR_API_KEY
-        with open(CONFIG_PATH, "w", encoding="utf-8") as f:
-            yaml.dump(config, f, default_flow_style=False, sort_keys=False)
-    except Exception:
-        pass
-
-
-def get_header_footer_zones(regions, norm_height=1000):
-    if not regions:
-        return None, None
-    y_tops, y_bottoms = [], []
-    for r in regions:
-        bbox = r.get("bbox_2d") if isinstance(r, dict) else getattr(r, "bbox_2d", None)
-        if bbox and len(bbox) >= 4:
-            y_tops.append(bbox[1])
-            y_bottoms.append(bbox[3])
-    if not y_tops:
-        return None, None
-    return min(y_tops) / norm_height, max(y_bottoms) / norm_height
-
-
-def extract_zone_text_pdf(pdf_path, page_num, y_start_frac, y_end_frac):
-    try:
-        import pymupdf as fitz
+def _infer_image(image_path: str) -> str:
+    """Run fine-tuned model on a single image file and return markdown string."""
+    import torch
+    from PIL import Image
 
-        doc = fitz.open(pdf_path)
-        page = doc[page_num]
-        h, w = page.rect.height, page.rect.width
-        rect = fitz.Rect(0, h * y_start_frac, w, h * y_end_frac)
-        text = page.get_text(clip=rect).strip()
-        doc.close()
-        return text
-    except Exception:
-        return ""
+    model, processor = _load_model()
 
+    img = Image.open(image_path).convert("RGB")
+    img = _resize_for_inference(img)
 
-def extract_pdf_text_in_band(pdf_path, page_num, y_start_frac, y_end_frac):
-    try:
-        import pymupdf as fitz
-
-        doc = fitz.open(pdf_path)
-        page = doc[page_num]
-        h = page.rect.height
-        y_lo = h * y_start_frac
-        y_hi = h * y_end_frac
-        words = page.get_text("words")
-        doc.close()
-        parts = []
-        for w in words:
-            if len(w) >= 5:
-                y0, y1 = float(w[1]), float(w[3])
-                if y0 < y_hi and y1 > y_lo:
-                    parts.append(w[4])
-        return " ".join(parts).strip()
-    except Exception:
-        return ""
-
-
-def ocr_zone(image_path, y_start_frac, y_end_frac):
-    zone_name = "header" if y_end_frac < 0.5 else "footer"
+    fd, resized_path = tempfile.mkstemp(suffix=".png")
+    os.close(fd)
     try:
-        from PIL import Image
+        img.save(resized_path, "PNG")
+
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image", "url": resized_path},
+                {"type": "text", "text": "Document Parsing:"},
+            ],
+        }]
+
+        inputs = processor.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(model.device)
+        inputs.pop("token_type_ids", None)
+
+        torch.cuda.empty_cache()
+
+        with torch.no_grad():
+            ids = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS)
+
+        result = processor.decode(
+            ids[0][inputs["input_ids"].shape[1]:],
+            skip_special_tokens=True,
+        )
+        return result.strip()
 
-        img = Image.open(image_path).convert("RGB")
-        w, h = img.size
-        y0 = max(0, int(h * y_start_frac))
-        y1 = min(h, int(h * y_end_frac))
-        if y1 <= y0:
-            return ""
-
-        crop = img.crop((0, y0, w, y1))
-        cw, ch = crop.size
-
-        if ch < MIN_CROP_HEIGHT or (cw * ch) < MIN_CROP_PIXELS:
-            need_h = max(ch, MIN_CROP_HEIGHT)
-            need_w = max(cw, 1)
-            if (need_w * need_h) < MIN_CROP_PIXELS:
-                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))
-            else:
-                canvas.paste(crop, (0, need_h - ch))
-            crop = canvas
-
-        fd, path = tempfile.mkstemp(suffix=".jpg")
-        os.close(fd)
+    finally:
         try:
-            crop.save(path, "JPEG", quality=ZONE_JPEG_QUALITY)
-            parser = get_parser()
-            out = parser.parse(path)
-            if not isinstance(out, list):
-                out = [out]
-            if out and getattr(out[0], "markdown_result", None):
-                return (out[0].markdown_result or "").strip()
-        finally:
-            try:
-                os.unlink(path)
-            except Exception:
-                pass
-    except Exception as e:
-        log.warning("[%s] ocr_zone failed: %s", zone_name, e, exc_info=True)
-    return ""
+            os.unlink(resized_path)
+        except Exception:
+            pass
 
 
 def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
@@ -219,7 +147,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)
@@ -231,12 +161,17 @@ 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
 
         uniq = uuid.uuid4().hex[:10]
-        img_path = os.path.join(tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{uniq}_{i}.png")
+        img_path = os.path.join(
+            tempfile.gettempdir(),
+            f"glmocr_page_{os.getpid()}_{uniq}_{i}.png",
+        )
         img.save(img_path, "PNG", compress_level=PAGE_PNG_COMPRESS_LEVEL)
         page_images.append(img_path)
         page_heights.append(img.height)
@@ -245,24 +180,6 @@ def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
     return page_images, page_heights
 
 
-def get_page_md_and_regions(page_result):
-    md = ""
-    if hasattr(page_result, "markdown_result") and page_result.markdown_result:
-        md = (page_result.markdown_result or "").strip()
-    regions = []
-    if hasattr(page_result, "json_result"):
-        jr = page_result.json_result
-        if isinstance(jr, dict) and "regions" in jr:
-            regions = jr.get("regions") or []
-        elif isinstance(jr, list) and len(jr) > 0:
-            r = jr[0] if isinstance(jr[0], list) else jr
-            if isinstance(r, list):
-                regions = r
-            elif isinstance(r, dict) and "regions" in r:
-                regions = r.get("regions") or []
-    return md, regions
-
-
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file."
@@ -271,92 +188,39 @@ def run_ocr(uploaded_file):
     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: List[int] = []
 
         if is_pdf:
-            page_images, page_heights = render_pdf_pages_to_images(path)
-            results = parser.parse(page_images)
+            page_images, _ = render_pdf_pages_to_images(path)
         else:
             page_images = [path]
-            page_heights = [1000]
-            results = parser.parse(path)
-
-        if not isinstance(results, list):
-            results = [results]
 
         all_pages = []
-        for page_num, page_result in enumerate(results):
-            page_md, regions = get_page_md_and_regions(page_result)
-            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)
-
-            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)
-
-            he = max(0.02, min(0.25, he))
-            fs = max(0.75, min(0.98, fs))
-
-            parts = []
-
-            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)
-            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(hdr.strip())
-
-            if page_md and page_md.strip():
-                parts.append(page_md.strip())
-
-            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():
-                    ftr_clean = ftr.strip()
-
-                    ftr_first_line = next(
-                        (ln.strip().lower() for ln in ftr_clean.splitlines() if ln.strip()),
-                        "",
-                    )
-                    already_present = ftr_first_line and any(
-                        ftr_first_line in part.lower() for part in parts
-                    )
-
-                    _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")
-                    _date_hits = len(_footer_date_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:
-                        parts.append(ftr_clean)
-
-            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)"
+        for page_num, img_path in enumerate(page_images):
+            log.info("Processing page %d / %d ...", page_num + 1, len(page_images))
+            page_md = _infer_image(img_path)
+            if page_md:
+                all_pages.append(page_md)
+
+        merged = (
+            "\n\n---page-separator---\n\n".join(all_pages)
+            if all_pages
+            else "(No content extracted)"
+        )
         return merged
 
     except Exception as e:
         import traceback
-
         log.exception("run_ocr failed: %s", e)
         return f"Error: {e}\n\n{traceback.format_exc()}"
 
     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
@@ -365,17 +229,17 @@ def run_ocr(uploaded_file):
 def _create_gradio_demo():
     import gradio as gr
 
-    with gr.Blocks(title="GLM-OCR (simple)") as demo:
-        gr.Markdown("# GLM-OCR")
+    with gr.Blocks(title="GLM-OCR Fine-tuned") as demo:
+        gr.Markdown("# GLM-OCR (Fine-tuned)")
         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 (markdown / HTML from model)")
+        out = gr.Textbox(lines=40, label="Output (markdown)")
         run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
     return demo
 
 
 if __name__ == "__main__":
-    _create_gradio_demo().launch()
+    _create_gradio_demo().launch()