app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
os.environ.setdefault("MOCR2_MAX_PIXELS", "1003520")

import spaces  # MUST be before any CUDA-touching import
import torch
import gradio as gr
import re
import json
import base64
import tempfile
from io import BytesIO
from pathlib import Path
from html import escape
from typing import Union, List, Optional, Tuple
from PIL import Image, ImageFile, ImageDraw

ImageFile.LOAD_TRUNCATED_IMAGES = True

from transformers import AutoModelForCausalLM, AutoProcessor

MODEL_ID = "zenosai/MonkeyOCRv2-B-Parsing"

# ── Prompts (from the official parsing pipeline) ──────────────────────────────

ALL_PROMPT = {
    "Caption": "Please output the text content from the image.",
    "List-item": "Please output the text content from the image.",
    "Page-footer": "Please output the text content from the image.",
    "Page-header": "Please output the text content from the image.",
    "Section-header": "Please output the text content from the image.",
    "Text": "Please output the text content from the image.",
    "Title": "Please output the text content from the image.",
    "Formula": "Please write out the expression of the formula in the image using LaTeX format.",
    "Table": "Please extract the table from the image and represent it in OTSL format.",
    "LAYOUT": "Please output the categories and coordinates of the document elements in reading order.",
}

# ── Model loading (module scope, eager .to("cuda")) ──────────────────────────

processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
    trust_remote_code=True,
    attn_implementation="sdpa",
).to("cuda").eval()

# ── Helper functions (ported from parse.py) ──────────────────────────────────


def _safe_eval(text: str):
    return eval(text, {"__builtins__": {}}, {})


def _normalize_item(item):
    if not isinstance(item, dict):
        return None
    if "bbox" not in item or "label" not in item:
        return None
    bbox = item["bbox"]
    label = item["label"]
    if not isinstance(bbox, (list, tuple)) or len(bbox) != 4:
        return None
    try:
        bbox = [float(v) for v in bbox]
    except Exception:
        return None
    if not isinstance(label, str):
        label = str(label)
    return {"bbox": bbox, "label": label}


def _normalize_list(obj):
    if not isinstance(obj, list):
        return []
    out = []
    for x in obj:
        nx = _normalize_item(x)
        if nx is not None:
            out.append(nx)
    return out


def _extract_balanced_blocks(text: str, lch: str, rch: str):
    res = []
    depth = 0
    start = -1
    for i, ch in enumerate(text):
        if ch == lch:
            if depth == 0:
                start = i
            depth += 1
        elif ch == rch and depth > 0:
            depth -= 1
            if depth == 0 and start != -1:
                res.append(text[start : i + 1])
                start = -1
    return res


def _dedup_keep_order(seq):
    seen = set()
    out = []
    for x in seq:
        if x not in seen:
            seen.add(x)
            out.append(x)
    return out


def _extract_tolerant_list_blocks(text: str):
    blocks = _extract_balanced_blocks(text, "[", "]")
    first = text.find("[")
    if first != -1:
        tail = text[first:].strip()
        if tail:
            lcnt, rcnt = tail.count("["), tail.count("]")
            if lcnt > rcnt:
                tail = tail + ("]" * (lcnt - rcnt))
            blocks.append(tail)
    return _dedup_keep_order(blocks)


def _extract_tolerant_dict_blocks(text: str):
    blocks = _extract_balanced_blocks(text, "{", "}")
    n = len(text)
    for i, ch in enumerate(text):
        if ch != "{":
            continue
        depth = 0
        end = None
        for j in range(i, n):
            cj = text[j]
            if cj == "{":
                depth += 1
            elif cj == "}":
                depth -= 1
                if depth == 0:
                    end = j + 1
                    break
        if end is None:
            blk = text[i:] + ("}" * max(depth, 1))
        else:
            blk = text[i:end]
        blocks.append(blk)
    return _dedup_keep_order(blocks)


def _parse_one_output(text: str):
    text = (text or "").strip()
    if not text:
        return []
    try:
        obj = _safe_eval(text)
        full = _normalize_list(obj)
        if full:
            return full
    except Exception:
        pass
    best = []
    for blk in _extract_tolerant_list_blocks(text):
        try:
            obj = _safe_eval(blk)
            cur = _normalize_list(obj)
            if len(cur) > len(best):
                best = cur
        except Exception:
            continue
    dict_items = []
    for blk in _extract_tolerant_dict_blocks(text):
        try:
            obj = _safe_eval(blk)
            nobj = _normalize_item(obj)
            if nobj is not None:
                dict_items.append(nobj)
        except Exception:
            continue
    if len(dict_items) > len(best):
        best = dict_items
    return best


def _map_bbox_to_image(bbox, w, h):
    x1, y1, x2, y2 = bbox
    x1 = x1 / 1000.0 * w
    x2 = x2 / 1000.0 * w
    y1 = y1 / 1000.0 * h
    y2 = y2 / 1000.0 * h
    if x1 > x2:
        x1, x2 = x2, x1
    if y1 > y2:
        y1, y2 = y2, y1
    x1 = max(0, min(int(round(x1)), w - 1 if w > 0 else 0))
    y1 = max(0, min(int(round(y1)), h - 1 if h > 0 else 0))
    x2 = max(x1 + 1, min(int(round(x2)), w))
    y2 = max(y1 + 1, min(int(round(y2)), h))
    return [x1, y1, x2, y2]


def otsl_to_html(otsl_str):
    if not otsl_str or not otsl_str.strip():
        return "<table></table>"
    rows_tokens = otsl_str.split("<nl>")
    if rows_tokens and rows_tokens[-1] == "":
        rows_tokens.pop()
    grid = []
    for r_idx, row_str in enumerate(rows_tokens):
        if not row_str.strip():
            if r_idx >= len(grid):
                grid.append([])
            continue
        parts = re.findall(r"<([a-z]+)>(.*?)(?=<[a-z]+>|$)", row_str)
        if r_idx >= len(grid):
            grid.append([])
        col_idx = 0
        for tag, content in parts:
            while True:
                while len(grid[r_idx]) <= col_idx:
                    grid[r_idx].append(None)
                if grid[r_idx][col_idx] is not None:
                    col_idx += 1
                else:
                    break
            if tag == "fcel" or tag == "ecel":
                text = content.strip() if tag == "fcel" else ""
                grid[r_idx][col_idx] = {
                    "text": text,
                    "rowspan": 1,
                    "colspan": 1,
                    "valid": True,
                }
                col_idx += 1
            elif tag == "lcel":
                search_c = col_idx - 1
                found = False
                while search_c >= 0:
                    if len(grid[r_idx]) > search_c:
                        cell = grid[r_idx][search_c]
                        if cell and cell.get("valid"):
                            cell["colspan"] += 1
                            found = True
                            break
                    search_c -= 1
                if found:
                    grid[r_idx][col_idx] = {"valid": False, "type": "lcel"}
                else:
                    grid[r_idx][col_idx] = {
                        "text": "",
                        "rowspan": 1,
                        "colspan": 1,
                        "valid": True,
                    }
                col_idx += 1
            elif tag == "ucel":
                search_r = r_idx - 1
                found = False
                while search_r >= 0:
                    if len(grid[search_r]) > col_idx:
                        cell = grid[search_r][col_idx]
                        if cell and cell.get("valid"):
                            cell["rowspan"] += 1
                            found = True
                            break
                    search_r -= 1
                if found:
                    grid[r_idx][col_idx] = {"valid": False, "type": "ucel"}
                else:
                    grid[r_idx][col_idx] = {
                        "text": "",
                        "rowspan": 1,
                        "colspan": 1,
                        "valid": True,
                    }
                col_idx += 1
            elif tag == "xcel":
                grid[r_idx][col_idx] = {"valid": False, "type": "xcel"}
                col_idx += 1
            else:
                col_idx += 1
    html_parts = ["<table>"]
    for row in grid:
        html_parts.append("<tr>")
        for cell in row:
            if cell is None:
                continue
            elif cell.get("valid"):
                attrs = []
                if cell["rowspan"] > 1:
                    attrs.append(f'rowspan="{cell["rowspan"]}"')
                if cell["colspan"] > 1:
                    attrs.append(f'colspan="{cell["colspan"]}"')
                attr_str = " " + " ".join(attrs) if attrs else ""
                text = escape(cell["text"])
                html_parts.append(f"<td{attr_str}>{text}</td>")
        html_parts.append("</tr>")
    html_parts.append("</table>")
    return "".join(html_parts)


def process_formula(content: str):
    content = content.strip("$").strip()
    content = re.sub(r"(?:\\quad\s*){5,}", r"\\quad ", content)
    content = re.sub(r"(?:\\qquad\s*){5,}", r"\\qquad ", content).strip()
    match = re.search(
        r"(?:\\quad|\\qquad|\\eqno)\s*\(([^()]*)\)\s*$"
        r"|\\tag\{([^{}]*)\}\s*$",
        content,
    )
    extracted = None
    if match:
        extracted = match.group(1)
        content = content[: match.start()].rstrip()
    begin_env = None
    has_end = False
    begin_match = re.match(r"^\\begin\{([^\}]+)\}", content)
    if begin_match:
        begin_env = begin_match.group(1)
        content = content[begin_match.end() :].lstrip()
        end_pattern = rf"\\end\{{{re.escape(begin_env)}\}}\s*$"
        end_match = re.search(end_pattern, content)
        if end_match:
            has_end = True
            content = content[: end_match.start()].rstrip()
    match = re.search(
        r"(?:\\quad|\\qquad|\\eqno)\s*\(([^()]*)\)\s*$"
        r"|\\tag\{([^{}]*)\}\s*$",
        content,
    )
    if match:
        extracted = match.group(1)
        content = content[: match.start()].rstrip()
    if begin_env:
        content = f"\\begin{{{begin_env}}}\n{content}\n\\end{{{begin_env}}}"
    return content, extracted


def detect_repeat_token(
    predicted_tokens: str,
    base_max_repeats: int = 4,
    window_size: int = 500,
    cut_from_end: int = 0,
    scaling_factor: float = 3.0,
):
    if cut_from_end > 0:
        predicted_tokens = predicted_tokens[:-cut_from_end]
    for seq_len in range(1, window_size // 2 + 1):
        candidate_seq = predicted_tokens[-seq_len:]
        max_repeats = int(base_max_repeats * (1 + scaling_factor / seq_len))
        repeat_count = 0
        pos = len(predicted_tokens) - seq_len
        if pos < 0:
            continue
        while pos >= 0:
            if predicted_tokens[pos : pos + seq_len] == candidate_seq:
                repeat_count += 1
                pos -= seq_len
            else:
                break
        if repeat_count > max_repeats:
            return True
    return False


def _should_retry_repeat_output(raw: str) -> bool:
    raw = raw or ""
    return detect_repeat_token(raw) or (
        len(raw) > 50 and detect_repeat_token(raw, cut_from_end=50)
    )


def image_to_png_data_uri(image: Image.Image) -> str:
    buffer = BytesIO()
    image.convert("RGB").save(buffer, format="PNG")
    encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
    return f"data:image/png;base64,{encoded}"


# ── Inference helpers ─────────────────────────────────────────────────────────


def _load_image_for_model(image: Image.Image, max_pixels: int = None, min_pixels: int = None) -> Image.Image:
    img = image.convert("RGB")
    if min_pixels and img.size[0] * img.size[1] < min_pixels:
        scale = (min_pixels / (img.size[0] * img.size[1])) ** 0.5
        new_size = (int(img.size[0] * scale), int(img.size[1] * scale))
        img = img.resize(new_size, Image.LANCZOS)
    if max_pixels and img.size[0] * img.size[1] > max_pixels:
        scale = (max_pixels / (img.size[0] * img.size[1])) ** 0.5
        new_size = (int(img.size[0] * scale), int(img.size[1] * scale))
        img = img.resize(new_size, Image.LANCZOS)
    return img


def _build_prompt(question: str) -> str:
    return (
        "system\nYou are a helpful assistant.\n"
        f"user\n<|vision_start|><|image_pad|><|vision_end|>"
        f"{question}\n"
        "assistant\n"
    )


def _generate(image: Image.Image, question: str, max_new_tokens: int = 5000, temperature: float = 0.0, do_sample: bool = False) -> str:
    """Single-image, single-question generation using transformers."""
    max_pixels = int(os.environ.get("MOCR2_MAX_PIXELS", "1003520"))
    img = _load_image_for_model(image, max_pixels=max_pixels)

    # Build messages and use apply_chat_template with tokenize=True
    messages = [
        {"role": "system", "content": [{"type": "text", "text": "You are a helpful assistant."}]},
        {"role": "user", "content": [
            {"type": "image", "image": img},
            {"type": "text", "text": question},
        ]},
    ]

    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to("cuda")

    # Remove keys not expected by model.generate (mm_token_type_ids from Qwen2VL processor)
    gen_inputs = {k: v for k, v in inputs.items() if k not in ("mm_token_type_ids",)}

    with torch.inference_mode():
        out = model.generate(
            **gen_inputs,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            temperature=temperature if do_sample else 1.0,
        )

    generated_ids = out[:, gen_inputs["input_ids"].shape[1]:]
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
    return text


def _batch_generate(images: List[Image.Image], questions: List[str], max_new_tokens: int = 5000) -> List[str]:
    """Batch generation with repeat-token retry."""
    results = []
    for img, q in zip(images, questions):
        raw = _generate(img, q, max_new_tokens=max_new_tokens)
        # Check for repeat tokens and retry with sampling
        if _should_retry_repeat_output(raw):
            for attempt in range(3):
                retry_temp = min(0.2 * (attempt + 1), 0.8)
                raw = _generate(img, q, max_new_tokens=max_new_tokens, temperature=retry_temp, do_sample=True)
                if not _should_retry_repeat_output(raw):
                    break
        results.append(raw)
    return results


# ── Two-stage parsing pipeline ────────────────────────────────────────────────


def get_layout(images: List[Image.Image]) -> List[List[dict]]:
    outputs = _batch_generate(images, [ALL_PROMPT["LAYOUT"]] * len(images), max_new_tokens=5000)
    page_layouts = []
    for i, out in enumerate(outputs):
        parsed = _parse_one_output(out)
        w, h = images[i].size
        mapped = []
        for item in parsed:
            mapped.append({
                "bbox": _map_bbox_to_image(item["bbox"], w, h),
                "label": item["label"],
            })
        page_layouts.append(mapped)
    return page_layouts


def parse_images(images: List[Image.Image]) -> Tuple[List[List[dict]], List[List[dict]]]:
    """Two-stage parsing: layout detection → element recognition."""
    layouts_per_page = get_layout(images)

    tasks = []
    for page_idx, items in enumerate(layouts_per_page):
        img = images[page_idx]
        w, h = img.size
        for item in items:
            x1, y1, x2, y2 = item["bbox"]
            x1 = max(0, min(x1, w - 1 if w > 0 else 0))
            y1 = max(0, min(y1, h - 1 if h > 0 else 0))
            x2 = max(x1 + 1, min(int(round(x2)), w))
            y2 = max(y1 + 1, min(int(round(y2)), h))
            label = item["label"]
            crop = img.crop((x1, y1, x2, y2))
            tasks.append({
                "image": crop,
                "bbox": [x1, y1, x2, y2],
                "label": label,
                "question": ALL_PROMPT.get(label, ""),
                "need_infer": label in ALL_PROMPT,
                "page_idx": page_idx,
            })

    infer_indices = [i for i, t in enumerate(tasks) if t["need_infer"]]
    infer_images = [tasks[i]["image"] for i in infer_indices]
    infer_questions = [tasks[i]["question"] for i in infer_indices]

    infer_outputs = _batch_generate(infer_images, infer_questions, max_new_tokens=5000) if infer_indices else []

    raw_outputs = [""] * len(tasks)
    for k, t_idx in enumerate(infer_indices):
        raw_outputs[t_idx] = infer_outputs[k]

    page_results = [[] for _ in images]
    for t, raw in zip(tasks, raw_outputs):
        label = t["label"]
        content = (raw or "").strip()
        if label == "Formula":
            content, extracted = process_formula(content)
            content = "$$\n" + content + "\n$$\n"
            if extracted:
                content = content + extracted
        elif label == "Table":
            content = otsl_to_html(content)
        elif label == "Picture":
            image_ref = image_to_png_data_uri(t["image"])
            content = f"![image]({image_ref})"
        elif label == "Title":
            content = "# " + content.replace("\n", "\n# ")
        elif label == "Section-header":
            content = "## " + content.replace("\n", "\n## ")
        elif not t["need_infer"]:
            content = ""
        rec = {
            "bbox": t["bbox"],
            "label": label,
            "content": content,
            "page_num": 1,
        }
        page_results[t["page_idx"]].append(rec)

    return page_results, layouts_per_page


def result_to_markdown(results: List[List[dict]], keep_header_footer: bool = False) -> str:
    lines = []
    for page_items in results:
        for item in page_items:
            if not keep_header_footer and item.get("label") in {"Page-header", "Page-footer"}:
                continue
            content = (item.get("content") or "").strip()
            if content:
                lines.append(content)
    md = "\n\n".join(lines).strip() + ("\n" if lines else "")
    md = md.replace("\ufffd", "")
    return md


def draw_layout(image: Image.Image, layout: List[dict]) -> Image.Image:
    """Draw bounding boxes on the image for layout visualization."""
    canvas = image.convert("RGB").copy()
    draw = ImageDraw.Draw(canvas)
    colors = {
        "Title": (255, 0, 0),
        "Section-header": (255, 128, 0),
        "Text": (0, 255, 0),
        "Formula": (0, 0, 255),
        "Table": (128, 0, 255),
        "Picture": (255, 0, 255),
        "Caption": (0, 255, 255),
        "Page-header": (128, 128, 128),
        "Page-footer": (128, 128, 128),
        "List-item": (255, 165, 0),
    }
    for i, it in enumerate(layout):
        x1, y1, x2, y2 = it["bbox"]
        label = it.get("label", "")
        color = colors.get(label, (255, 0, 0))
        draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
        ty = max(0, y1 - 12)
        draw.text((x1, ty), f"{i}: {label}", fill=color)
    return canvas


# ── Gradio UI ─────────────────────────────────────────────────────────────────

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""


def _estimate_duration(image, *args, **kwargs):
    return 120


@spaces.GPU(duration=_estimate_duration)
def parse_document(
    image: Image.Image,
    keep_header_footer: bool = False,
    show_layout: bool = True,
) -> Tuple[str, Optional[Image.Image]]:
    """Parse a document image into structured Markdown.

    Uses MonkeyOCRv2-B-Parsing, a visual-text foundation model for Document AI.
    The model performs two-stage parsing: (1) layout detection to identify
    document elements (text, tables, formulas, images) and their reading order,
    then (2) content recognition for each element.

    Args:
        image: Document image to parse.
        keep_header_footer: Whether to keep page headers/footers in the markdown.
        show_layout: Whether to visualize detected layout bounding boxes.

    Returns:
        A tuple of (parsed_markdown, layout_visualization_image).
    """
    if image is None:
        return "Please upload a document image.", None

    images = [image]
    results, layouts = parse_images(images)
    markdown = result_to_markdown(results, keep_header_footer=keep_header_footer)
    layout_img = draw_layout(image, layouts[0]) if show_layout and layouts else None

    return markdown, layout_img


with gr.Blocks() as demo:
    gr.Markdown(
        """
        # MonkeyOCRv2: Document AI Parsing Demo

        Upload a document image (receipts, papers, forms, tables, formulas, etc.) and the model will
        detect the layout, extract text, tables (as HTML), formulas (as LaTeX), and pictures —
        returning structured Markdown.

        **Model**: [`zenosai/MonkeyOCRv2-B-Parsing`](https://huggingface.co/zenosai/MonkeyOCRv2-B-Parsing)
        | **Paper**: [MonkeyOCRv2](https://huggingface.co/papers/2607.11562)
        | **GitHub**: [Yuliang-Liu/MonkeyOCRv2](https://github.com/Yuliang-Liu/MonkeyOCRv2)
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(
                label="Document Image",
                type="pil",
                height=500,
            )
            with gr.Accordion("Options", open=False):
                keep_hf = gr.Checkbox(
                    label="Keep headers & footers",
                    value=False,
                )
                show_layout_cb = gr.Checkbox(
                    label="Show layout visualization",
                    value=True,
                )
            run_btn = gr.Button("Parse Document", variant="primary")

        with gr.Column(scale=1):
            markdown_out = gr.Code(
                label="Parsed Markdown",
                language="markdown",
                lines=25,
            )
            layout_out = gr.Image(
                label="Layout Visualization",
                type="pil",
                height=400,
            )

    run_btn.click(
        fn=parse_document,
        inputs=[input_image, keep_hf, show_layout_cb],
        outputs=[markdown_out, layout_out],
        api_name="parse",
    )

    gr.Examples(
        examples=[
            ["examples/en.JPEG"],
            ["examples/exampaper.jpg"],
            ["examples/table.png"],
            ["examples/formula.png"],
        ],
        inputs=[input_image],
        outputs=[markdown_out, layout_out],
        fn=parse_document,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)