Add single-region mode, code fence languages, measured GPU duration

README.md

@@ -35,6 +35,11 @@ The app follows the authors' two-stage pipeline:
    converted to HTML, equations to LaTeX, using the authors' post-processors
    (vendored under `NaviOCR/`).
 
+A third **Single region** mode skips layout and runs the authors' one-block path
+(`block_parse`) over the whole image with the prompt for a chosen block type —
+table, formula, code, seal, or a chart / scientific figure, which the model
+converts into the table it implies.
+
 Outputs: rendered document, Markdown source, layout overlay, and the raw block
 list as JSON.
 

app.py

@@ -11,12 +11,16 @@ Two-stage pipeline, faithful to the authors' reference implementation
    de-rotated, and recognized with the block-type-specific prompt and sampling
    parameters from `NaviOCR/vlm_utils/NaviOCR_client.py`, then post-processed
    (OTSL tables -> HTML, LaTeX equation fixes) with the authors' post-processors.
+
+A third mode skips layout and runs the authors' single-region path
+(`NaviOCRClient.block_parse`) on the whole image, which is how the model card
+demonstrates chart-to-table extraction, seal reading and table/formula crops.
 """
 
 import base64
 import io
-import json
 import os
+import re
 import tempfile
 import time
 from dataclasses import asdict
@@ -38,6 +42,7 @@ from NaviOCR.vlm_utils.NaviOCR_client import (
     NaviOCRClient,
 )
 from NaviOCR.vlm_utils.post_process.otsl2html import convert_otsl_to_html
+from NaviOCR.vlm_utils.structs import ContentBlock
 from NaviOCR.vlm_utils.vlm_client import SamplingParams
 
 MODEL_ID = "StarDoc-AI/NaviDC-OCR"
@@ -90,6 +95,18 @@ BLOCK_COLORS = {
 }
 DEFAULT_COLOR = (117, 117, 117)
 
+# Block types the single-region mode exposes, with the authors' prompt keys.
+REGION_TASKS = [
+    ("Text", "text"),
+    ("Table \u2192 HTML", "table"),
+    ("Formula \u2192 LaTeX", "formula"),
+    ("Code", "code"),
+    ("Chart / scientific figure \u2192 table", "char"),
+    ("Seal", "seal"),
+]
+# The model prefixes recognized code with its own language marker, e.g. `<_Python_>`.
+CODE_LANG_RE = re.compile(r"^\s*<_([A-Za-z0-9+#._\- ]+)_>\s*")
+
 
 def _sampling_params(task: str, max_new_tokens: int) -> SamplingParams:
     """Authors' per-task sampling params, with a bounded generation length."""
@@ -165,6 +182,15 @@ def draw_layout(image: Image.Image, blocks: list) -> Image.Image:
     return Image.alpha_composite(canvas.convert("RGBA"), overlay).convert("RGB")
 
 
+def _fenced_code(content: str) -> str:
+    match = CODE_LANG_RE.match(content)
+    language = ""
+    if match:
+        language = match.group(1).strip().lower().replace(" ", "")
+        content = content[match.end() :]
+    return f"```{language}\n{content}\n```"
+
+
 def blocks_to_markdown(image: Image.Image, blocks: list, drop_paratext: bool):
     """Assemble reading-ordered blocks into Markdown (raw + display variants)."""
     parts: list[str] = []
@@ -193,7 +219,7 @@ def blocks_to_markdown(image: Image.Image, blocks: list, drop_paratext: bool):
         elif block_type == "char":
             parts.append(convert_otsl_to_html(content) or content)
         elif block_type in {"code", "algorithm"}:
-            parts.append(f"```\n{content}\n```")
+            parts.append(_fenced_code(content))
         elif block_type in CAPTION_TYPES:
             parts.append(f"*{content}*")
         elif block_type == "seal":
@@ -219,10 +245,20 @@ def _write_markdown(markdown: str) -> str:
     return path
 
 
-@spaces.GPU(duration=150)
+def _estimate_duration(*args, **kwargs) -> int:
+    """A dense page measured 67 s end-to-end; scale mildly with the token cap."""
+    max_new_tokens = 2048
+    if len(args) > 4:
+        max_new_tokens = args[4]
+    max_new_tokens = int(kwargs.get("max_new_tokens", max_new_tokens) or 2048)
+    return int(min(180, 55 + 0.02 * max_new_tokens))
+
+
+@spaces.GPU(duration=_estimate_duration)
 def parse_document(
     image: Image.Image,
     layout_mode: str = "Detection",
+    region_task: str = "text",
     drop_paratext: bool = True,
     max_new_tokens: int = 2048,
     progress=gr.Progress(track_tqdm=True),
@@ -232,8 +268,11 @@ def parse_document(
     Args:
         image: A document page — a digital page, a scan, or a camera photo.
         layout_mode: "Detection" for axis-aligned boxes (digital pages, flat
-            scans) or "Segmentation" for multi-point polygons (camera-captured,
-            curved or crumpled pages).
+            scans), "Segmentation" for multi-point polygons (camera-captured,
+            curved or crumpled pages), or "Region" to skip layout and recognize
+            the whole image as one block.
+        region_task: The block type used in "Region" mode — one of text, table,
+            formula, code, char (chart/scientific figure), seal.
         drop_paratext: Drop headers, footers, page numbers and margin notes.
         max_new_tokens: Generation cap per region.
 
@@ -247,7 +286,33 @@ def parse_document(
     started = time.time()
     page = image.convert("RGB") if isinstance(image, Image.Image) else Image.open(image).convert("RGB")
     helper = client.helper
-    mode = layout_mode if layout_mode in LAYOUT_PROMPTS else "Detection"
+    mode = layout_mode if layout_mode in LAYOUT_PROMPTS else "Region"
+
+    # ---- single-region mode: the authors' block_parse path ----------------
+    if mode == "Region":
+        task = region_task if region_task in DEFAULT_PROMPTS else "text"
+        crop = helper.resize_by_need(page)
+        output = client.client.predict(
+            crop,
+            DEFAULT_PROMPTS[task],
+            _sampling_params(task, max_new_tokens),
+        )
+        block = ContentBlock(type=task, bbox=[[0.0, 0.0], [1.0, 1.0]], content=output)
+        blocks = helper.post_process([block]) or [block]
+        raw_markdown, display_markdown = blocks_to_markdown(page, blocks, False)
+        seconds = time.time() - started
+        report = (
+            f"Single region recognized as `{task}` \u2014 {seconds:.1f}s.  \n"
+            f"Switch to a full-page mode to run layout analysis first."
+        )
+        return (
+            page,
+            display_markdown,
+            raw_markdown,
+            [dict(item) for item in blocks],
+            _write_markdown(raw_markdown),
+            report,
+        )
 
     # ---- stage 1: layout ------------------------------------------------
     layout_image = helper.prepare_for_layout(page)  # resized to 1036x1036
@@ -342,11 +407,21 @@ off flat scans **and** photographed / crumpled pages, and returns Markdown.
                 image = gr.Image(label="Document page", type="pil", height=460)
                 layout_mode = gr.Radio(
                     choices=[
-                        ("Boxes — digital pages & flat scans", "Detection"),
-                        ("Multi-point — photos, curved or crumpled pages", "Segmentation"),
+                        ("Full page, boxes — digital pages & flat scans", "Detection"),
+                        (
+                            "Full page, multi-point — photos, curved or crumpled pages",
+                            "Segmentation",
+                        ),
+                        ("Single region — the image is one table / formula / …", "Region"),
                     ],
                     value="Detection",
-                    label="Layout mode",
+                    label="Parsing mode",
+                )
+                region_task = gr.Dropdown(
+                    choices=REGION_TASKS,
+                    value="table",
+                    label="Region type",
+                    visible=False,
                 )
                 run_button = gr.Button("Parse document", variant="primary")
                 report = gr.Markdown()
@@ -378,14 +453,14 @@ off flat scans **and** photographed / crumpled pages, and returns Markdown.
 
         gr.Examples(
             examples=[
-                ["examples/journal_page.jpg", "Detection"],
-                ["examples/crumpled_page.jpg", "Segmentation"],
-                ["examples/table.png", "Detection"],
-                ["examples/formula.png", "Detection"],
-                ["examples/code.png", "Detection"],
-                ["examples/scientific_figure.png", "Detection"],
+                ["examples/journal_page.jpg", "Detection", "table"],
+                ["examples/crumpled_page.jpg", "Segmentation", "table"],
+                ["examples/table.png", "Region", "table"],
+                ["examples/formula.png", "Region", "formula"],
+                ["examples/code.png", "Region", "code"],
+                ["examples/scientific_figure.png", "Region", "char"],
             ],
-            inputs=[image, layout_mode],
+            inputs=[image, layout_mode, region_task],
             outputs=[overlay, document, markdown_source, blocks_json, markdown_file, report],
             fn=parse_document,
             cache_examples=True,
@@ -393,10 +468,18 @@ off flat scans **and** photographed / crumpled pages, and returns Markdown.
             label="Examples from the NaviDC-OCR model card",
         )
 
+    layout_mode.change(
+        fn=lambda mode: gr.update(visible=(mode == "Region")),
+        inputs=[layout_mode],
+        outputs=[region_task],
+        show_api=False,
+        queue=False,
+    )
+
     gr.on(
         triggers=[run_button.click],
         fn=parse_document,
-        inputs=[image, layout_mode, drop_paratext, max_new_tokens],
+        inputs=[image, layout_mode, region_task, drop_paratext, max_new_tokens],
         outputs=[overlay, document, markdown_source, blocks_json, markdown_file, report],
     )