Add spatial HTML preview mode using grounding boxes

app.py

@@ -9,6 +9,7 @@ import shutil
 from PIL import Image, ImageDraw, ImageFont, ImageOps
 import fitz
 import re
+import ast
 import numpy as np
 import base64
 import html as html_lib
@@ -148,10 +149,87 @@ math[display="block"] { display: block; overflow-x: auto; max-width: 100%; }
 .math-preview pre code { background: none; padding: 0; }
 .math-preview blockquote { border-left: 4px solid #ccc; margin: 0.8em 0; padding: 0.4em 1em; color: #555; background: #fafafa; }
 .math-preview img { max-width: 100%; height: auto; display: block; margin: 0.8em 0; }
+.math-preview .ocr-gap, .mathjax-preview .ocr-gap { width: 100%; }
 .math-fallback { color: #888; font-style: italic; }
 </style>
+<script>
+(() => {
+  if (window.__ocrMathJaxInit) return;
+  window.__ocrMathJaxInit = true;
+
+  if (!window.MathJax) {
+    window.MathJax = {
+      tex: {
+        inlineMath: [['\\\\(', '\\\\)'], ['$', '$']],
+        displayMath: [['\\\\[', '\\\\]'], ['$$', '$$']]
+      },
+      options: {
+        skipHtmlTags: ['script', 'noscript', 'style', 'textarea', 'pre', 'code']
+      }
+    };
+  }
+
+  const typeset = () => {
+    if (window.MathJax?.typesetPromise) {
+      const nodes = Array.from(document.querySelectorAll('.mathjax-preview, .spatial-preview'));
+      if (nodes.length) window.MathJax.typesetPromise(nodes).catch(() => {});
+    }
+  };
+  window.__typesetOcrMath = typeset;
+
+  const ensureScript = () => {
+    if (document.getElementById('mathjax-ocr-preview')) return;
+    const script = document.createElement('script');
+    script.id = 'mathjax-ocr-preview';
+    script.async = true;
+    script.src = 'https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js';
+    script.onload = () => setTimeout(typeset, 20);
+    document.head.appendChild(script);
+  };
+
+  ensureScript();
+  setTimeout(typeset, 100);
+
+  const observer = new MutationObserver((mutations) => {
+    for (const m of mutations) {
+      for (const n of m.addedNodes) {
+        if (n.nodeType !== 1) continue;
+        if (n.matches?.('.mathjax-preview, .spatial-preview') || n.querySelector?.('.mathjax-preview, .spatial-preview')) {
+          setTimeout(typeset, 30);
+          return;
+        }
+      }
+    }
+  });
+  observer.observe(document.body, { childList: true, subtree: true });
+})();
+</script>
 """
 
+def _inject_spatial_gap_placeholders(text: str):
+    """Preserve runs of blank lines so OCR spacing is visible in preview."""
+    gaps: dict[str, int] = {}
+    counter = [0]
+
+    def repl(m):
+        key = f'ZZOCRGAP{counter[0]}ZZ'
+        counter[0] += 1
+        # Two newlines are a normal paragraph break; extras represent vertical spacing.
+        gaps[key] = max(1, len(m.group(0)) - 2)
+        return f'\n\n{key}\n\n'
+
+    return re.sub(r'\n{3,}', repl, text), gaps
+
+def _restore_spatial_gap_placeholders(html: str, gaps: dict[str, int]) -> str:
+    if not gaps:
+        return html
+    for key, extra_lines in gaps.items():
+        gap_em = min(10.0, 0.9 * extra_lines)
+        block = f'<div class="ocr-gap" style="height:{gap_em:.2f}em"></div>'
+        html = html.replace(f'<p>{key}</p>', block)
+        html = html.replace(key, block)
+    return html
+
 def _to_mathml(latex: str, display: bool) -> str:
     """Convert a LaTeX string to MathML. Falls back to a code block on error."""
     # Fix OCR error: \frac{n/m} (single-argument fraction) → \frac{n}{m}
@@ -203,6 +281,7 @@ def to_math_html(text: str) -> str:
     text = re.sub(r'\\\[.*', '', text, flags=re.DOTALL)
     # Replace inline math \(...\) with placeholder tokens
     text = re.sub(r'\\\((.+?)\\\)', inline_math, text)
+    text, gaps = _inject_spatial_gap_placeholders(text)
 
     # Run markdown on text that now contains only safe placeholder tokens
     html = md_lib.markdown(text, extensions=['tables', 'fenced_code', 'sane_lists', 'nl2br'])
@@ -233,63 +312,170 @@ def to_math_html(text: str) -> str:
     for key, literal in literals.items():
         html = html.replace(key, html_lib.escape(literal))
 
+    html = _restore_spatial_gap_placeholders(html, gaps)
+
     return f'<div class="math-preview">{html}</div>'
 
 def to_mathjax_html(text: str) -> str:
     """Render markdown to HTML and typeset math client-side with MathJax."""
     if not text:
         return ""
+    text, gaps = _inject_spatial_gap_placeholders(text)
     html = md_lib.markdown(text, extensions=['tables', 'fenced_code', 'sane_lists', 'nl2br'])
-    return f"""<!doctype html>
-<html>
-<head>
-  <meta charset="utf-8" />
-  <style>
-    body {{
-      margin: 0;
-      font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
-      font-size: 15px;
-      line-height: 1.8;
-      color: #1a1a1a;
-    }}
-    .mathjax-preview {{
-      padding: 1.5em;
-      max-width: 100%;
-      overflow-x: auto;
-    }}
-    .mathjax-preview h1 {{ font-size: 1.8em; font-weight: 700; margin: 1em 0 0.4em; border-bottom: 2px solid #e0e0e0; padding-bottom: 0.3em; }}
-    .mathjax-preview h2 {{ font-size: 1.4em; font-weight: 600; margin: 1em 0 0.4em; border-bottom: 1px solid #e0e0e0; padding-bottom: 0.2em; }}
-    .mathjax-preview h3 {{ font-size: 1.15em; font-weight: 600; margin: 0.9em 0 0.3em; }}
-    .mathjax-preview p  {{ margin: 0.6em 0; }}
-    .mathjax-preview ul, .mathjax-preview ol {{ padding-left: 1.8em; margin: 0.5em 0; }}
-    .mathjax-preview li {{ margin: 0.25em 0; }}
-    .mathjax-preview table {{ border-collapse: collapse; width: 100%; margin: 1em 0; font-size: 0.95em; }}
-    .mathjax-preview th, .mathjax-preview td {{ border: 1px solid #ccc; padding: 0.45em 0.75em; text-align: left; }}
-    .mathjax-preview th {{ background: #f2f2f2; font-weight: 600; }}
-    .mathjax-preview tr:nth-child(even) {{ background: #fafafa; }}
-    .mathjax-preview code {{ background: #f4f4f4; padding: 0.15em 0.4em; border-radius: 3px; font-family: 'Courier New', monospace; font-size: 0.88em; }}
-    .mathjax-preview pre  {{ background: #f4f4f4; padding: 1em; border-radius: 5px; overflow-x: auto; margin: 0.8em 0; }}
-    .mathjax-preview pre code {{ background: none; padding: 0; }}
-    .mathjax-preview blockquote {{ border-left: 4px solid #ccc; margin: 0.8em 0; padding: 0.4em 1em; color: #555; background: #fafafa; }}
-    .mathjax-preview img {{ max-width: 100%; height: auto; display: block; margin: 0.8em 0; }}
-  </style>
-  <script>
-    window.MathJax = {{
-      tex: {{
-        inlineMath: [['\\\\(', '\\\\)'], ['$', '$']],
-        displayMath: [['\\\\[', '\\\\]'], ['$$', '$$']]
-      }},
-      options: {{
-        skipHtmlTags: ['script', 'noscript', 'style', 'textarea', 'pre', 'code']
-      }}
-    }};
-  </script>
-  <script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
-</head>
-<body>
-  <div class="mathjax-preview">{html}</div>
-</body>
-</html>"""
+    html = _restore_spatial_gap_placeholders(html, gaps)
+    return f'<div class="mathjax-preview">{html}</div>'
+
+def _grounding_blocks_from_raw(raw_text: str):
+    if not raw_text:
+        return []
+
+    pattern = re.compile(r'<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>', re.DOTALL)
+    blocks = []
+    last_end = 0
+
+    for m in pattern.finditer(raw_text):
+        label = m.group(1).strip() or "text"
+        coord_text = m.group(2).strip()
+        text_chunk = raw_text[last_end:m.start()].strip()
+        last_end = m.end()
+
+        try:
+            coords = ast.literal_eval(coord_text)
+        except (SyntaxError, ValueError):
+            continue
+
+        if isinstance(coords, (tuple, list)) and coords and isinstance(coords[0], (int, float)):
+            coords = [coords]
+        if not isinstance(coords, list):
+            continue
+
+        boxes = [c for c in coords if isinstance(c, (list, tuple)) and len(c) >= 4]
+        if not boxes:
+            continue
+
+        x1 = max(0.0, min(float(c[0]) for c in boxes))
+        y1 = max(0.0, min(float(c[1]) for c in boxes))
+        x2 = min(999.0, max(float(c[2]) for c in boxes))
+        y2 = min(999.0, max(float(c[3]) for c in boxes))
+        if x2 <= x1 or y2 <= y1:
+            continue
+
+        blocks.append({
+            "label": label,
+            "text": text_chunk,
+            "x1": x1,
+            "y1": y1,
+            "x2": x2,
+            "y2": y2,
+        })
+
+    return blocks
+
+def to_spatial_html(raw_text: str, markdown_text: str) -> str:
+    """Render OCR content using grounding boxes for spatially-positioned blocks."""
+    blocks = _grounding_blocks_from_raw(raw_text)
+    if not blocks:
+        return to_mathjax_html(markdown_text)
+
+    used_text = 0
+    rendered = []
+    palette = {
+        "title": "#8b5cf6",
+        "text": "#2563eb",
+        "image": "#059669",
+        "table": "#d97706",
+        "formula": "#dc2626",
+    }
+
+    for i, b in enumerate(sorted(blocks, key=lambda x: (x["y1"], x["x1"]))):
+        label = b["label"]
+        color = palette.get(label.lower(), "#4b5563")
+        body = b["text"].strip()
+        if body:
+            used_text += len(body)
+            body_text, gaps = _inject_spatial_gap_placeholders(body)
+            body_html = md_lib.markdown(body_text, extensions=['tables', 'fenced_code', 'sane_lists', 'nl2br'])
+            body_html = _restore_spatial_gap_placeholders(body_html, gaps)
+        else:
+            body_html = ""
+
+        if not body_html:
+            body_html = f"<p><em>{html_lib.escape(label)}</em></p>"
+
+        left = b["x1"] / 999.0 * 100.0
+        top = b["y1"] / 999.0 * 100.0
+        width = max(1.0, (b["x2"] - b["x1"]) / 999.0 * 100.0)
+        height = max(1.2, (b["y2"] - b["y1"]) / 999.0 * 100.0)
+
+        rendered.append(
+            f"""
+<article class="spatial-block" style="left:{left:.2f}%; top:{top:.2f}%; width:{width:.2f}%; min-height:{height:.2f}%; --block-color:{color};">
+  <header>{html_lib.escape(label)}</header>
+  <section>{body_html}</section>
+</article>
+"""
+        )
+
+    fallback = ""
+    if markdown_text and used_text < max(120, int(len(markdown_text) * 0.4)):
+        fallback_html = to_mathjax_html(markdown_text)
+        fallback = f"""
+<details class="spatial-fallback">
+  <summary>Show full linear markdown rendering</summary>
+  {fallback_html}
+</details>
+"""
+
+    return f"""
+<style>
+.spatial-preview {{
+  padding: 1rem;
+}}
+.spatial-canvas {{
+  position: relative;
+  width: 100%;
+  min-height: 72vh;
+  aspect-ratio: 1 / 1.35;
+  background: linear-gradient(180deg, #fcfdff 0%, #f7f9fc 100%);
+  border: 1px solid #d8dee9;
+  border-radius: 8px;
+  overflow: auto;
+}}
+.spatial-block {{
+  position: absolute;
+  box-sizing: border-box;
+  border: 1px solid var(--block-color);
+  background: color-mix(in srgb, var(--block-color) 7%, white);
+  border-radius: 6px;
+  padding: 0.35rem 0.5rem;
+  overflow: hidden;
+}}
+.spatial-block > header {{
+  font-size: 11px;
+  font-weight: 700;
+  letter-spacing: 0.03em;
+  text-transform: uppercase;
+  color: var(--block-color);
+  margin-bottom: 0.25rem;
+}}
+.spatial-block > section {{
+  font-size: 13px;
+  line-height: 1.35;
+}}
+.spatial-block p {{ margin: 0.2rem 0; }}
+.spatial-fallback {{
+  margin-top: 1rem;
+  padding-top: 0.5rem;
+  border-top: 1px solid #d8dee9;
+}}
+</style>
+<div class="spatial-preview mathjax-preview">
+  <div class="spatial-canvas">
+    {''.join(rendered)}
+  </div>
+  {fallback}
+</div>
+"""
 
 def embed_images(markdown, crops):
     if not crops:
@@ -461,6 +647,9 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
                 with gr.Tab("HTML + MathJax", id="tab_html"):
                     html_out = gr.HTML("")
                     html_source_out = gr.Code(label="Generated HTML Source", language="html", lines=16)
+                with gr.Tab("Spatial HTML", id="tab_spatial"):
+                    spatial_out = gr.HTML("")
+                    spatial_source_out = gr.Code(label="Spatial HTML Source", language="html", lines=16)
                 with gr.Tab("Boxes", id="tab_boxes"):
                     img_out = gr.Image(type="pil", height=500, show_label=False)
                 with gr.Tab("Cropped Images", id="tab_crops"):
@@ -520,7 +709,7 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
         elif image is not None:
             cleaned, markdown, raw, img_out, crops = process_image(image, task, custom_prompt)
         else:
-            return "Error: Upload a file or image", "", "", "", "", None, [], gr.DownloadButton(visible=False)
+            return "Error: Upload a file or image", "", "", "", "", "", "", None, [], gr.DownloadButton(visible=False)
 
         # Text tab: convert \[...\] → $$...$$ and \(...\) → $...$ for readability
         text_display = re.sub(r'\\\[(.+?)\\\]',
@@ -534,12 +723,13 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
         dl_tmp.close()
 
         mathjax_html = to_mathjax_html(markdown)
+        spatial_html = to_spatial_html(raw, markdown)
 
-        return (text_display, to_math_html(markdown), mathjax_html, mathjax_html, raw, img_out, crops,
+        return (text_display, to_math_html(markdown), mathjax_html, mathjax_html, spatial_html, spatial_html, raw, img_out, crops,
                 gr.DownloadButton(value=dl_tmp.name, visible=True))
 
     submit_event = btn.click(run, [input_img, file_in, task, prompt, page_selector],
-                             [text_out, md_out, html_out, html_source_out, raw_out, img_out, gallery, download_btn])
+                             [text_out, md_out, html_out, html_source_out, spatial_out, spatial_source_out, raw_out, img_out, gallery, download_btn])
     submit_event.then(select_boxes, [task], [tabs])
 
 if __name__ == "__main__":