Spaces:

PaddlePaddle
/

hpd-parsing

Running on Zero

Render output as Markdown + bbox visualization (instead of raw text)

by WEISHU - opened 6 days ago

PaddlePaddle org 6 days ago

Motivation

parse_document() currently dumps the raw output of model.generate_hpd(...) (including <BLOCK>/<FORK>/<CHILD>/[bbox] control tags) straight into a gr.Textbox. This has two downsides:

Users see a tagged raw string rather than a readable document structure — formulas and tables aren't rendered at all.
The model actually emits a layout bounding box for every region ([x1, y1, x2, y2], normalized to 0-1000), but it's completely discarded — users have no way to visually sanity-check whether the model's region segmentation is accurate.

Changes

Added hpd_postprocess.py. Without touching the generate_hpd call itself, it does two things with the raw output:

parse_blocks(raw_text): parses the <BLOCK> <type> [bbox] <CHILD> <content> stream and keeps each block's type, bbox (0-1000 normalized coordinates), and cleaned text. The formula-cleaning pipeline is reused as-is from the repo's own eval/hpd_to_markdown.py (simplify_left_right / clean_formula_tail / normalize_arith, unmodified, just copied over). Note: the official script's parsing logic discards [bbox] entirely, so a separate bbox-preserving parser was written rather than reusing that function directly.
draw_boxes_on_image(image, blocks): draws the boxes on the original image, color-coded by block type, for visual verification.

Changes to app.py:

parse_document() now returns (boxed_image, markdown_text) instead of the raw string.
The output column is switched from a single gr.Textbox to gr.Image (bbox visualization) + gr.Markdown (rendered parse result).
import spaces / @spaces.GPU is wrapped in try/except so it gracefully degrades to a no-op decorator when the spaces package isn't available (useful for local/offline debugging). Behavior on the actual ZeroGPU Space is unchanged.

Diff

diff --git a/app.py b/app.py
index 8a88cb9..c675271 100644
--- a/app.py
+++ b/app.py
@@ -4,12 +4,23 @@ import os
 
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
-import spaces  # MUST be before any torch/CUDA import
+try:
+    import spaces  # MUST be before any torch/CUDA import; only present on HF ZeroGPU
+    GPU_DECORATOR = spaces.GPU(duration=60)
+except ImportError:
+    def GPU_DECORATOR(fn):
+        return fn
+
 import torch
 import gradio as gr
 from transformers import AutoModel, AutoTokenizer
 
-MODEL_ID = "PaddlePaddle/HPD-Parsing"
+from hpd_postprocess import parse_blocks, blocks_to_markdown, draw_boxes_on_image
+
+MODEL_ID = os.environ.get(
+    "HPD_MODEL_PATH",
+    "PaddlePaddle/HPD-Parsing",
+)
 
 # --- Image preprocessing (mirrors the repo's image_preprocess.py) ----------
 import torchvision.transforms as T
@@ -123,7 +134,7 @@ print("Model loaded.")
 
 # --- Inference ---------------------------------------------------------------
 
-@spaces.GPU(duration=60)
+@GPU_DECORATOR
 def parse_document(image, use_fork, use_mtp, max_new_tokens):
     """Parse a document image into structured text using HPD-Parsing.
 
@@ -132,9 +143,14 @@ def parse_document(image, use_fork, use_mtp, max_new_tokens):
         use_fork: Whether to use hierarchical parallel decoding (fork path).
         use_mtp: Whether to use P-MTP speculative decoding for the parent branch.
         max_new_tokens: Maximum number of new tokens to generate.
+
+    Returns:
+        A tuple ``(boxed_image, markdown_text)`` where ``boxed_image`` is the
+        input image annotated with typed bounding boxes and ``markdown_text``
+        is the cleaned, tag-free markdown reconstruction of the parse.
     """
     if image is None:
-        return "Please upload a document image first."
+        return None, "Please upload a document image first."
 
     pixel_values = load_image_from_pil(image).to(torch.bfloat16).to("cuda")
 
@@ -149,7 +165,11 @@ def parse_document(image, use_fork, use_mtp, max_new_tokens):
         num_speculative_tokens=6,
         batch_children=False,
     )
-    return response
+
+    blocks = parse_blocks(response)
+    markdown_text = blocks_to_markdown(blocks) or response
+    boxed_image = draw_boxes_on_image(image, blocks)
+    return boxed_image, markdown_text
 
 
 # --- Gradio UI ---------------------------------------------------------------
@@ -180,7 +200,8 @@ with gr.Blocks() as demo:
                 max_tokens = gr.Slider(label="Max new tokens", minimum=512, maximum=16000, value=8000, step=256)
 
         with gr.Column(scale=1):
-            output_text = gr.Textbox(label="Parsed Output", lines=25)
+            output_image = gr.Image(label="Detected Layout (bounding boxes)")
+            output_markdown = gr.Markdown(label="Parsed Output (rendered Markdown)")
 
     gr.Examples(
         examples=[
@@ -188,7 +209,7 @@ with gr.Blocks() as demo:
             ["sample_invoice.png", True, True, 8000],
         ],
         inputs=[input_image, use_fork_cb, use_mtp_cb, max_tokens],
-        outputs=output_text,
+        outputs=[output_image, output_markdown],
         fn=parse_document,
         cache_examples=True,
         cache_mode="lazy",
@@ -197,8 +218,8 @@ with gr.Blocks() as demo:
     run_btn.click(
         fn=parse_document,
         inputs=[input_image, use_fork_cb, use_mtp_cb, max_tokens],
-        outputs=output_text,
+        outputs=[output_image, output_markdown],
         api_name="parse_document",
     )
 
-demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)
+demo.launch(mcp_server=True)

Full content of new file: `hpd_postprocess.py`

"""Post-processing utilities for HPD-Parsing demo: turns the raw
``<BLOCK> <type> [bbox] <CHILD> <content>`` stream into (a) clean markdown text
and (b) a list of typed bounding boxes for visualization.

Formula-cleaning helpers (``simplify_left_right`` / ``clean_formula_tail`` /
``normalize_arith``) are copied verbatim from the model repo's
``eval/hpd_to_markdown.py`` so the produced markdown matches the official
OmniDocBench post-processing.
"""

import re

from PIL import ImageDraw, ImageFont

# --- Formula-cleaning helpers (copied from eval/hpd_to_markdown.py) --------

_TALL = re.compile(
    r'\\d?frac|\\tfrac|\\cfrac|\\binom|\\sqrt'
    r'|\\sum|\\prod|\\coprod|\\int|\\iint|\\iiint|\\oint'
    r'|\\bigcup|\\bigcap|\\bigoplus|\\bigotimes|\\bigsqcup'
    r'|\\begin\{'
    r'|\\overbrace|\\underbrace|\\overset|\\underset|\\stackrel'
    r'|\\substack|\\atop|\\\\'
)


def _scan_delims(s):
    out = []
    for m in re.finditer(r'\\(left|right)\s*', s):
        dm = re.match(r'\\[a-zA-Z]+|\\.|.', s[m.end():])
        if not dm:
            continue
        out.append({'kind': m.group(1), 'delim': dm.group(0),
                    'start': m.start(), 'end': m.end() + dm.end()})
    return out


def simplify_left_right(s: str) -> str:
    """Downgrade `\\left( ... \\right)` with no tall inner structure to plain `( )`."""
    if '\\left' not in s:
        return s
    stack, pairs = [], []
    for d in _scan_delims(s):
        if d['kind'] == 'left':
            stack.append(d)
        elif stack:
            pairs.append((stack.pop(), d))
    edits = []
    for L, R in pairs:
        if L['delim'] == '(' and R['delim'] == ')' and not _TALL.search(s[L['end']:R['start']]):
            edits.append((L['start'], L['end'], '('))
            edits.append((R['start'], R['end'], ')'))
    for st, en, rep in sorted(edits, key=lambda x: x[0], reverse=True):
        s = s[:st] + rep + s[en:]
    return s


_ELLIPSIS = r'(?:\\dots|\\cdots|\\ldots|\\dotsb|\\dotsc)'
_CLOSER = r'(?:\\right\s*[.\}\]\)]|\\end\s*\{(?:array|matrix|cases|bmatrix|pmatrix|vmatrix|smallmatrix)\})'
_TAIL_WRAP = re.compile(r'^(?P<core>.*?)(?P<wrap>\s*(?:\\\]|\\\)|\$\$))?\s*$', re.DOTALL)


def clean_formula_tail(s: str) -> str:
    """Strip degenerate formula tails (repeated/dangling ellipses, stray `\\quad`)."""
    if not s:
        return s
    m = _TAIL_WRAP.match(s)
    core, wrap = m.group('core'), m.group('wrap') or ''
    prev = None
    while prev != core:
        prev = core
        core = re.sub(r'(' + _ELLIPSIS + r')(?:\s*' + _ELLIPSIS + r')+', r'\1', core)
        core = re.sub(r'(?P<keep>' + _CLOSER + r')\s*(?:\\q?quad\s*)*' + _ELLIPSIS + r'\s*$',
                      lambda mm: mm.group('keep'), core)
        core = re.sub(r'(?:\s*\\q?quad)+\s*' + _ELLIPSIS + r'\s*$', '', core)
        core = re.sub(r'(?:\s*\\q?quad)+\s*$', '', core)
        core = core.rstrip()
    return core + wrap


_OP_MAP = {
    '≈': r'\approx', '≠': r'\neq', '≤': r'\leq', '≥': r'\geq', '×': r'\times',
    '÷': r'\div', '±': r'\pm', '∓': r'\mp', '·': r'\cdot', '∙': r'\cdot',
    '⋅': r'\cdot', '∗': '*', '−': '-', '≡': r'\equiv', '∝': r'\propto',
    '∞': r'\infty', '√': r'\sqrt', '→': r'\to', '≪': r'\ll', '≫': r'\gg',
}
_ARITH_ALLOWED = re.compile(r'^[0-9A-Za-z\s=+\-*/^_().,:;<>|%!\u4e00-\u9fff' + ''.join(_OP_MAP.keys()) + r']+$')
_ARITH_HASOP = re.compile(r'[=+\-*/' + ''.join(_OP_MAP.keys()) + r']')
_KNOWN_FUNCS = {'sin', 'cos', 'tan', 'cot', 'sec', 'csc', 'log', 'ln', 'exp',
                'lim', 'max', 'min', 'det', 'mod', 'arcsin', 'arccos', 'arctan', 'sqrt'}
_CJK_RUN = re.compile(r'[\u4e00-\u9fff]+')
_MATH_SPAN = re.compile(r'(\\\[.*?\\\]|\$\$.*?\$\$|\\\(.*?\\\)|\$.*?\$)', re.DOTALL)

WRAP_CJK_IN_ARITH = True


def _convert_unicode_ops(s: str) -> str:
    for k, v in _OP_MAP.items():
        s = s.replace(k, (v + ' ') if v.startswith('\\') else v)
    if WRAP_CJK_IN_ARITH:
        s = _CJK_RUN.sub(lambda m: r'\text{' + m.group(0) + '}', s)
    return re.sub(r'[ \t]{2,}', ' ', s)


def _is_pure_arith_line(line: str) -> bool:
    t = line.strip()
    if not t or '\\(' in t or '\\[' in t or '$' in t or '<' in t:
        return False
    if not WRAP_CJK_IN_ARITH and re.search(r'[\u4e00-\u9fff]', t):
        return False
    if not _ARITH_ALLOWED.match(t) or not _ARITH_HASOP.search(t):
        return False
    return all(w.lower() in _KNOWN_FUNCS for w in re.findall(r'[A-Za-z]{2,}', t))


def normalize_arith(text: str) -> str:
    """Normalize Unicode operators to LaTeX and wrap pure-arithmetic lines as `\\( .. \\)`."""
    if not text:
        return text
    text = _MATH_SPAN.sub(lambda m: _convert_unicode_ops(m.group(0)), text)
    out = []
    for line in text.split('\n'):
        if _is_pure_arith_line(line):
            out.append('\\( ' + _convert_unicode_ops(line.strip()) + ' \\)')
        else:
            out.append(line)
    return '\n'.join(out)


def clean_text(text: str, simplify_left_paren=True, clean_formula_tail_flag=True,
               norm_formula_flag=True) -> str:
    """Apply the same per-block cleaning steps as the official hpd_to_markdown.py."""
    text = text.strip()
    text = text.replace('The image is too blurry to recognize any text content.', '').strip()
    text = text.replace(
        "The image contains no text or characters. It is a graphical element (a horizontal "
        "line with a vertical line) and does not contain any chart, graph, or data points "
        "that can be extracted. Therefore, the correct OCR output is an empty string.", ""
    ).strip()
    if not text or text == '[Non-Text]':
        return ''
    if text.startswith('\\[') and not text.endswith('\n\\]'):
        text += '\n\\]'
    if text.startswith('<table>') and not text.endswith('</table>'):
        text += '</table>'
    if '\\[\n' in text and '\\\\' not in text:
        text = text.replace('\\[\n', '\\(').replace('\n\\]', '\\)')
    text = text.replace('\\)  $^{'},^{'}$ \n\n\\(')
    if '÷' in text and '\\(' not in text:
        text = '\\( ' + text + ' \\)'
    text = re.sub(r'\\tag\s*\{[^{}]*\}', '', text)
    text = text.replace('\\supset', '\\sqsupset')
    if simplify_left_paren:
        text = simplify_left_right(text)
    if clean_formula_tail_flag:
        text = clean_formula_tail(text)
    if norm_formula_flag:
        text = normalize_arith(text)
    return text


# --- Block parsing (bbox-preserving, unlike the official script) -----------

# type + [bbox], usually followed by <FORK|CHILD|BLOCK>. Container blocks such
# as `list [x1,y1,x2,y2]` may have no trailing tag at all -- after splitting on
# `<BLOCK>` the header is simply the entire segment -- so the tag is optional.
_BLOCK_HEADER = re.compile(
    r'([a-zA-Z_]+)\s*\[\s*([-\d.,\s]+)\]\s*(?:<(?:FORK|CHILD|BLOCK)>)?'
)
_NO_CONTENT_TYPES = {'chart', 'seal'}


def parse_blocks(raw_text: str):
    """Parse the ``<BLOCK> <type> [bbox] <CHILD> <content>`` stream.

    Returns a list of dicts: ``{"type": str, "bbox": [x1,y1,x2,y2] | None, "text": str}``.
    ``bbox`` values are in the model's native 0-1000 normalized coordinate space
    (confirmed via real inference: max observed values ~922/589 against a
    1240x1754 source image). ``text`` is the cleaned markdown for that block;
    it is empty for container blocks (e.g. ``list``, ``table`` wrapper) and for
    ``chart``/``seal`` blocks, matching the official markdown-export behavior.
    """
    blocks = []
    segments = raw_text.split('<BLOCK>')[1:]
    for seg in segments:
        header_m = _BLOCK_HEADER.match(seg.strip())
        # segments always start right after a <BLOCK>, so the header should be
        # at the very start; fall back to a bare type token if bbox is absent.
        type_m = re.match(r'\s*([a-zA-Z_]+)', seg)
        block_type = type_m.group(1) if type_m else 'unknown'

        bbox = None
        if header_m:
            block_type = header_m.group(1)
            nums = [float(x) for x in re.split(r'[,\s]+', header_m.group(2).strip()) if x]
            if len(nums) == 4:
                bbox = nums

        text = ''
        if block_type.lower() not in _NO_CONTENT_TYPES:
            content_m = re.search(r'<CHILD>(.*)', seg, re.DOTALL)
            if content_m:
                # stop at the next control tag if any leaked through
                raw_content = re.split(r'<(?:FORK|CHILD|BLOCK)>', content_m.group(1))[0]
                text = clean_text(raw_content)

        blocks.append({'type': block_type, 'bbox': bbox, 'text': text})
    return blocks


def blocks_to_markdown(blocks) -> str:
    """Join the non-empty block texts in reading order into one markdown string."""
    return '\n\n'.join(b['text'] for b in blocks if b['text']).strip()


# --- Bounding-box visualization ---------------------------------------------

_TYPE_COLORS = {
    'header': '#e74c3c',
    'title': '#e67e22',
    'text': '#2980b9',
    'list': '#8e44ad',
    'table': '#27ae60',
    'table_caption': '#16a085',
    'figure': '#2c3e50',
    'figure_caption': '#34495e',
    'chart': '#f39c12',
    'seal': '#c0392b',
    'formula': '#d35400',
}
_DEFAULT_COLOR = '#7f8c8d'


def draw_boxes_on_image(image, blocks):
    """Draw typed bounding boxes on a copy of ``image``.

    ``blocks`` bbox values are 0-1000 normalized coordinates in ``[x1, y1, x2, y2]``
    order; they are rescaled to ``image``'s actual pixel size before drawing.
    Blocks with ``bbox is None`` are skipped (still counted, just not drawn).
    """
    if image is None:
        return None
    canvas = image.convert('RGB').copy()
    draw = ImageDraw.Draw(canvas)
    width, height = canvas.size

    try:
        font = ImageFont.load_default()
    except Exception:
        font = None

    for block in blocks:
        bbox = block.get('bbox')
        if not bbox:
            continue
        x1, y1, x2, y2 = bbox
        px1, py1 = x1 / 1000.0 * width, y1 / 1000.0 * height
        px2, py2 = x2 / 1000.0 * width, y2 / 1000.0 * height
        color = _TYPE_COLORS.get(block['type'].lower(), _DEFAULT_COLOR)
        draw.rectangle([px1, py1, px2, py2], outline=color, width=2)
        label = block['type']
        text_y = max(0, py1 - 12)
        if font is not None:
            draw.text((px1 + 1, text_y), label, fill=color, font=font)
        else:
            draw.text((px1 + 1, text_y), label, fill=color)

    return canvas

AsikHamim94

6 days ago

Add a HTML Preview side by side. Also adding model config or params will be great.

Upload images, audio, and videos by dragging in the text input, pasting, or clicking here.

Tap or paste here to upload images

· Sign up or log in to comment

Render output as Markdown + bbox visualization (instead of raw text)

Motivation

Changes

Diff

Full content of new file: hpd_postprocess.py

Full content of new file: `hpd_postprocess.py`