app.py

#!/usr/bin/env python3

import logging
import os
import re
import tempfile
import uuid
from typing import List, Tuple

log = logging.getLogger("glmocr_simple_app")
logging.basicConfig(level=logging.INFO)

# ── Fine-tuned model repo on HuggingFace ─────────────────────────────────────
MERGED_MODEL_DIR = os.environ.get("MODEL_DIR", "SimpleCodeAI/glm-ocr-finetuned")

RENDER_SCALE = 2.0
PAD_LEFT_FRAC = 0.035
PAD_RIGHT_FRAC = 0.10
PAD_TOP_FRAC = 0.018
PAD_BOTTOM_FRAC = 0.018

ENABLE_CONTRAST = True
CONTRAST_FACTOR = 1.18
ENABLE_UNSHARP = True
UNSHARP_RADIUS = 0.78
UNSHARP_PERCENT = 76
UNSHARP_THRESHOLD = 1

PAGE_PNG_COMPRESS_LEVEL = 3
MAX_IMAGE_SIDE = 1568
MAX_NEW_TOKENS = 3000

# ── 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):
    from PIL import ImageEnhance, ImageFilter

    if ENABLE_CONTRAST:
        img = ImageEnhance.Contrast(img).enhance(CONTRAST_FACTOR)
    if ENABLE_UNSHARP:
        img = img.filter(
            ImageFilter.UnsharpMask(
                radius=UNSHARP_RADIUS,
                percent=UNSHARP_PERCENT,
                threshold=UNSHARP_THRESHOLD,
            )
        )
    return img


def _resize_for_inference(img):
    """Resize image preserving aspect ratio so longest side <= MAX_IMAGE_SIDE."""
    from PIL import Image

    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)


def _build_table(rows: List[str]) -> str:
    parsed = []

    for row in rows:
        parts = re.split(r"\s{2,}", row.strip())
        if len(parts) >= 3:
            date = parts[0]
            amount = parts[-1] if re.search(r"\d+\.\d{2}", parts[-1]) else ""
            desc = " ".join(parts[1:-1]) if amount else " ".join(parts[1:])
            parsed.append((date, desc, amount))

    if not parsed:
        return "\n".join(rows)

    table = [
        "| POSTING DATE | DESCRIPTION | AMOUNT |",
        "| :--- | :--- | ---: |",
    ]
    for date, desc, amount in parsed:
        table.append(f"| {date} | {desc} | {amount} |")
    return "\n".join(table)


def _normalize_transactions(text: str) -> str:
    """
    Convert loose transaction lines into structured markdown tables.
    """
    lines = text.split("\n")
    result: List[str] = []
    table_buffer: List[str] = []
    in_section = False

    for line in lines:
        if re.search(r"POSTING DATE.*DESCRIPTION.*AMOUNT", line, re.IGNORECASE):
            in_section = True
            table_buffer = []
            continue

        if in_section and (line.strip() == "" or line.strip().startswith("Subtotal")):
            if table_buffer:
                result.append(_build_table(table_buffer))
                table_buffer = []
            in_section = False
            result.append(line)
            continue

        if in_section:
            table_buffer.append(line)
        else:
            result.append(line)

    if table_buffer:
        result.append(_build_table(table_buffer))

    return "\n".join(result)


def _fix_missing_amounts(text: str) -> str:
    lines = text.split("\n")
    fixed: List[str] = []

    for line in lines:
        if re.search(r"\d{2}/\d{2}", line) and not re.search(r"\d+\.\d{2}", line):
            match = re.search(r"(\d{1,3}(?:,\d{3})*\.\d{2})$", line)
            if match:
                line += f" {match.group(1)}"
        fixed.append(line)

    return "\n".join(fixed)


def _html_to_markdown_tables(text: str) -> str:
    try:
        from bs4 import BeautifulSoup
    except Exception:
        return text

    soup = BeautifulSoup(text, "html.parser")

    for table in soup.find_all("table"):
        rows = []
        for tr in table.find_all("tr"):
            cols = [td.get_text(strip=True) for td in tr.find_all(["td", "th"])]
            if cols:
                rows.append(cols)

        if rows:
            md = []
            header = rows[0]
            md.append("| " + " | ".join(header) + " |")
            md.append("| " + " | ".join(["---"] * len(header)) + " |")

            for row in rows[1:]:
                padded = row + [""] * (len(header) - len(row))
                md.append("| " + " | ".join(padded[: len(header)]) + " |")

            table.replace_with("\n".join(md))

    return str(soup)


def _clean_markdown(text: str) -> str:
    """Post-process markdown to fix table formatting and section structure."""
    text = _html_to_markdown_tables(text)
    text = _fix_missing_amounts(text)
    text = _normalize_transactions(text)

    lines = text.split('\n')
    cleaned = []
    in_table = False

    for line in lines:
        if '|' in line and line.strip().startswith('|'):
            if not in_table:
                if cleaned and cleaned[-1].strip():
                    cleaned.append('')
                in_table = True

            line = re.sub(r'\s*\|\s*', ' | ', line)
            line = re.sub(r'\s+', ' ', line)
            cleaned.append(line.strip())
        elif in_table and line.strip() == '':
            in_table = False
            cleaned.append('')
        else:
            in_table = False
            if line.strip() or (cleaned and cleaned[-1].strip()):
                cleaned.append(line)

    text = '\n'.join(cleaned)
    text = re.sub(r'(#{1,6})\s*([^\n]+)', r'\1 \2', text)
    text = re.sub(r'\n([•\-\*])\s+', r'\n\1 ', text)
    text = '\n'.join(line.rstrip() for line in text.split('\n'))
    text = re.sub(r'\n{4,}', '\n\n\n', text)
    return text.strip()


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

    model, processor = _load_model()

    img = Image.open(image_path).convert("RGB")
    img = _resize_for_inference(img)

    fd, resized_path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    try:
        img.save(resized_path, "PNG")

        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "url": resized_path},
                {
                    "type": "text",
                    "text": (
                        "Document Parsing to markdown.\n"
                        "Rules:\n"
                        "1) Preserve rows exactly in reading order.\n"
                        "2) Keep transaction blocks as markdown tables.\n"
                        "3) Do not invent rows or sample/template data.\n"
                        "4) Keep right-most amount values."
                    ),
                },
            ],
        }]

        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)

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        with torch.no_grad():
            ids = model.generate(
                **inputs,
                max_new_tokens=MAX_NEW_TOKENS,
                do_sample=False,
                repetition_penalty=1.1,
            )

        result = processor.decode(
            ids[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        )
        
        # Clean up the markdown output
        return _clean_markdown(result.strip())

    finally:
        try:
            os.unlink(resized_path)
        except Exception:
            pass


def render_pdf_pages_to_images(pdf_path: str) -> Tuple[List[str], List[int]]:
    import pymupdf as fitz
    from PIL import Image

    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)
        img = _enhance_raster_for_ocr(img)

        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

        uniq = uuid.uuid4().hex[:10]
        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)

    doc.close()
    return page_images, page_heights


def run_ocr(uploaded_file):
    if uploaded_file is None:
        return "Please upload a 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")

        if is_pdf:
            page_images, _ = render_pdf_pages_to_images(path)
        else:
            page_images = [path]

        all_pages = []
        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)
                ):
                    os.unlink(p)
            except Exception:
                pass


def _create_gradio_demo():
    import gradio as gr

    with gr.Blocks(title="GLM-OCR Fine-tuned") as demo:
        gr.Markdown("# GLM-OCR (Fine-tuned)")
        gr.Markdown("Upload a PDF or image to extract structured markdown content.")
        
        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)"
        )
        run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
        
    return demo


if __name__ == "__main__":
    _load_model()
    _create_gradio_demo().launch()