| """Single-file OvisOCR3 document parsing and cross-page table reconstruction. |
| |
| Reading guide: public entry point -> document pipeline -> text/table helpers -> |
| model backend -> browser renderer. Import the text helpers without installing |
| vLLM or browser dependencies; those are loaded only when used. |
| |
| Algorithm: parse pages independently, then visit adjacent table-bearing pages |
| from last to first. Write each merged table into the earlier page and render |
| that updated page before its next prediction. |
| Keep non-table source untouched. |
| |
| Support: the document pipeline accepts one merged table per adjacent pair and |
| binds the earlier page's last table to the later page's first table. Multiple |
| merges need explicit fragment bindings via apply_pair(); nested or incomplete |
| HTML tables are rejected on writeback. This is a conservative binding policy, |
| not a general table-fragment matching algorithm. |
| |
| Page numbers are 1-based; table fragment indices are 0-based. Model predictions |
| use pair-local pages 1 and 2; saved Markdown and provenance use document pages. |
| Raw predictions, intermediate files and partial results are retained. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import base64 |
| import copy |
| import html |
| import io |
| import json |
| import math |
| import re |
| import time |
| import traceback |
| import uuid |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
| from typing import Callable, Mapping, NamedTuple |
|
|
| NEGATIVE = "No cross-page table exists between these pages." |
| TABLE = re.compile(r"<table\b[^>]*>.*?</table\s*>", re.I | re.S) |
| IMAGE_REF = re.compile( |
| r"(?<![\w/])(?:page_(\d+)/)?images/bbox_(\d+)_(\d+)_(\d+)_(\d+)\.jpg" |
| ) |
| REF = re.compile(r"page_(\d+)/images/") |
| CATEGORIES = set( |
| "abstract algorithm aside_text chart content display_formula doc_title figure_title footer footer_image footnote formula_number header header_image image inline_formula number paragraph_title reference_content seal table text vertical_text vision_footnote".split() |
| ) |
|
|
| FragmentId = tuple[int, int] |
| Backend = Callable[[list[str], str, Path, int], tuple[str, str]] |
|
|
|
|
| class GenerationResult(NamedTuple): |
| """Tuple-compatible raw model output; cleaning belongs to the page adapter.""" |
|
|
| text: str |
| finish_reason: str |
|
|
|
|
| class RenderResult(NamedTuple): |
| """Tuple-compatible render result; aliases map rendered crops to originals.""" |
|
|
| image_path: str |
| rewritten_source: str |
| alias_to_original: dict[str, str] |
| diagnostics: dict |
|
|
|
|
| @dataclass |
| class PageState: |
| """One document page and the original fragments behind its current tables.""" |
|
|
| number: int |
| image_path: str |
| original_markdown: str |
| markdown: str |
| finish_reason: str |
| parse_diagnostics: dict | None = None |
| table_sources: list[set[FragmentId]] = field(init=False) |
|
|
| def __post_init__(self): |
| self.table_sources = [ |
| {(self.number, index)} |
| for index, _ in enumerate(TABLE.finditer(self.markdown)) |
| ] |
|
|
|
|
| @dataclass |
| class PairInput: |
| """Two model inputs with explicit document/pair-local coordinate mapping.""" |
|
|
| page_numbers: tuple[int, int] |
| image_paths: list[str] |
| markdowns: list[str] |
| alias_to_original: dict[str, str] |
| rendered_pages: set[int] |
|
|
|
|
| @dataclass |
| class DocumentState: |
| pages: list[PageState] = field(default_factory=list) |
| chunks: list[list[int]] = field(default_factory=list) |
| steps: list[dict] = field(default_factory=list) |
| merged_tables: list[dict] = field(default_factory=list) |
| parse_warnings: list[dict] = field(default_factory=list) |
|
|
|
|
| |
|
|
|
|
| def reconstruct_document( |
| image_paths, |
| output, |
| *, |
| model_path=".", |
| stage1_prompt="document_structured", |
| mathjax=None, |
| backend: Backend | None = None, |
| renderer=None, |
| ): |
| """Reconstruct ordered page images into a document, keeping all evidence. |
| |
| Use a fresh output directory. Stage 1 accepts document_markdown[.txt] or |
| document_structured[.txt]; prompts are loaded from model_path/prompts. |
| Returns the saved result.json dictionary. Truncated generations, recovered |
| JSON or invalid structured fields return PARTIAL with issue-specific warnings; |
| other processing failures save partial output and raise RuntimeError. |
| |
| For research, inject backend(paths, prompt, folder, stage) -> (raw, reason) |
| and/or renderer.render(source, owner, paths, folder, name) -> RenderResult. |
| Injected renderers remain caller-owned. Defaults load vLLM and Chromium lazily. |
| Engine settings, prompt bytes and renderer resize policy are unchanged. |
| """ |
| stage1_prompt = stage1_prompt.removesuffix(".txt") |
| if stage1_prompt not in {"document_markdown", "document_structured"}: |
| raise ValueError("Choose document_markdown or document_structured") |
| images = _document_images(image_paths) |
| if Path(output).exists(): |
| raise FileExistsError("Use a fresh output directory") |
|
|
| prompt_dir = Path(model_path) / "prompts" |
| page_prompt = (prompt_dir / f"{stage1_prompt}.txt").read_text().removesuffix("\n") |
| table_prompt = (prompt_dir / "crosspage_table.txt").read_text().removesuffix("\n") |
| if backend is None: |
| backend = VllmBackend(model_path) |
| owns_renderer = renderer is None |
| if owns_renderer: |
| renderer = _BrowserRenderer( |
| mathjax or Path(__file__).resolve().parent / "assets/mathjax.js" |
| ) |
| try: |
| result = _run_document( |
| images, |
| output, |
| backend, |
| renderer, |
| page_prompt=page_prompt, |
| table_prompt=table_prompt, |
| stage1_prompt=stage1_prompt, |
| ) |
| finally: |
| if owns_renderer: |
| renderer.close() |
|
|
| if result["error"] and not result["error"].get("recoverable", False): |
| raise RuntimeError(f"Document incomplete; inspect {output}/failure.json") |
| return result |
|
|
|
|
| def _run_document( |
| image_paths, output, backend, renderer, *, page_prompt, table_prompt, stage1_prompt |
| ): |
| """Parse pages, reconstruct adjacent pairs backwards, then persist results. |
| |
| Recovery is confined to Stage 1 output parsing. A failed window stops further |
| reconstruction while preserving earlier pages, merges and failure details. |
| No benchmark labels are used. The output directory must not already exist. |
| """ |
| images = _document_images(image_paths) |
| output = Path(output) |
| output.mkdir(parents=True, exist_ok=False) |
| save( |
| output / "inputs.json", |
| { |
| "images": images, |
| "stage1_prompt": stage1_prompt, |
| "page_prompt": page_prompt, |
| "table_prompt": table_prompt, |
| }, |
| ) |
| state = DocumentState() |
| error = None |
| context = {"stage": 1, "pages": []} |
| try: |
| for number, image in enumerate(images, 1): |
| context = {"stage": 1, "pages": [number]} |
| page = _parse_page( |
| number, image, output, backend, page_prompt, stage1_prompt |
| ) |
| state.pages.append(page) |
| warning = _page_warning(page) |
| if warning: |
| state.parse_warnings.append(warning) |
|
|
| state.chunks = find_chunks([page.markdown for page in state.pages]) |
| save( |
| output / "chunks.json", |
| { |
| "chunks": state.chunks, |
| "basis": "Stage 1 predicted HTML tables only", |
| }, |
| ) |
| for left, right in _adjacent_pairs(state.chunks): |
| context = {"stage": 2, "pages": [left, right]} |
| folder = output / "stage2" / f"pair_{left:04d}_{right:04d}" |
| folder.mkdir(parents=True) |
| pair = _prepare_pair( |
| state.pages[left - 1 : right], images, folder, renderer |
| ) |
| tables = _predict_pair(pair, folder, backend, table_prompt) |
| _apply_merge(state, pair, tables) |
| state.steps.append( |
| { |
| "pages": list(pair.page_numbers), |
| "merged": bool(tables), |
| "rendered_pages": sorted(pair.rendered_pages), |
| } |
| ) |
| save( |
| folder / "result.json", |
| { |
| "tables": tables, |
| "pages": [ |
| state.pages[number - 1].markdown for number in pair.page_numbers |
| ], |
| "groups": copy.deepcopy(state.merged_tables), |
| }, |
| ) |
| except Exception as exc: |
| error = { |
| "type": type(exc).__name__, |
| "message": str(exc), |
| "recoverable": False, |
| **context, |
| } |
| save(output / "failure.json", {**error, "traceback": traceback.format_exc()}) |
| return _finalize_document(state, images, output, error) |
|
|
|
|
| def _parse_page(number, image, output, backend, page_prompt, stage1_prompt): |
| """Retain raw output, apply the existing cleanup, then adapt the page format.""" |
| folder = output / "stage1" / f"page_{number:04d}" |
| folder.mkdir(parents=True) |
| generation = GenerationResult(*backend([image], page_prompt, folder, 1)) |
| raw, reason = generation |
| (folder / "raw.txt").write_text(raw) |
| cleaned = _clean_truncated_repeats(raw.strip()) |
| save( |
| folder / "generation.json", |
| { |
| "finish_reason": reason, |
| "complete": reason == "stop", |
| "tail_changed": cleaned != raw.strip(), |
| "raw_chars": len(raw.strip()), |
| "cleaned_chars": len(cleaned), |
| }, |
| ) |
| (folder / "cleaned_raw.txt").write_text(cleaned) |
| if reason not in ("stop", "length"): |
| raise ValueError(f"Stage 1 page {number}: finish_reason={reason}") |
|
|
| diagnostics = None |
| if stage1_prompt == "document_structured": |
| projected, blocks, diagnostics = structured_to_markdown(cleaned) |
| save(folder / "blocks.json", blocks) |
| save(folder / "adapter.json", diagnostics) |
| (folder / "projected.md").write_text(projected) |
| markdown = page_markdown(projected, number) if projected.strip() else "" |
| else: |
| markdown = page_markdown(cleaned, number) |
| (folder / "page.md").write_text(markdown) |
| return PageState(number, image, markdown, markdown, reason, diagnostics) |
|
|
|
|
| def _page_warning(page): |
| """Separate generation truncation, JSON recovery and structured schema issues.""" |
| codes = [] |
| if page.finish_reason == "length": |
| codes.append("generation_truncated") |
| diagnostics = page.parse_diagnostics |
| if diagnostics is not None: |
| if not diagnostics["json_complete"]: |
| codes.append("json_incomplete") |
| if diagnostics["schema_issues"]: |
| codes.append("schema_invalid") |
| if ( |
| diagnostics["invalid_content_objects"] |
| or not diagnostics["emitted_contents"] |
| ): |
| codes.append("content_invalid") |
| if not codes: |
| return None |
| warning = {"page": page.number, "finish_reason": page.finish_reason, "codes": codes} |
| if diagnostics is not None: |
| warning["diagnostics"] = diagnostics |
| return warning |
|
|
|
|
| def _adjacent_pairs(chunks): |
| """Yield (2, 3), then (1, 2) for a three-page table-bearing run.""" |
| for chunk in chunks: |
| for left in range(chunk[-1] - 1, chunk[0] - 1, -1): |
| yield left, left + 1 |
|
|
|
|
| def _prepare_pair(pages, images, folder, renderer): |
| page_numbers = tuple(page.number for page in pages) |
| document_to_local = {number: local for local, number in enumerate(page_numbers, 1)} |
| paths, markdowns, aliases, rendered_pages = [], [], {}, set() |
| for page in pages: |
| path, markdown = page.image_path, page.markdown |
| if markdown != page.original_markdown: |
| rendered = RenderResult( |
| *renderer.render( |
| markdown, page.number, images, folder, f"page_{page.number:04d}" |
| ) |
| ) |
| path, markdown = rendered.image_path, rendered.rewritten_source |
| aliases.update(rendered.alias_to_original) |
| rendered_pages.add(page.number) |
| paths.append(path) |
| markdowns.append(remap(markdown, document_to_local)) |
| return PairInput(page_numbers, paths, markdowns, aliases, rendered_pages) |
|
|
|
|
| def _predict_pair(pair, folder, backend, table_prompt): |
| page_context = "\n\n".join( |
| f'<INPUT_PAGE index="{index}">\n{markdown}\n</INPUT_PAGE>' |
| for index, markdown in enumerate(pair.markdowns, 1) |
| ) |
| prompt = table_prompt.removesuffix("\n") + "\n\nInput Markdown:\n\n" + page_context |
| if all(TABLE.search(markdown) for markdown in pair.markdowns): |
| generation = GenerationResult(*backend(pair.image_paths, prompt, folder, 2)) |
| else: |
| generation = GenerationResult( |
| json.dumps([{"source_pages": [], "html": NEGATIVE}]), "stop" |
| ) |
| save( |
| folder / "gate_skip.json", |
| {"reason": "README table gate", "model_called": False}, |
| ) |
| (folder / "raw.txt").write_text(generation.text) |
| save(folder / "generation.json", {"finish_reason": generation.finish_reason}) |
| if generation.finish_reason != "stop": |
| raise ValueError( |
| f"Stage 2 pair {list(pair.page_numbers)}: finish_reason={generation.finish_reason}" |
| ) |
| tables = parse_tables(generation.text) |
| if len(tables) > 1: |
| raise ValueError( |
| "Multiple merges need explicit fragment bindings; output retained." |
| ) |
| return [ |
| dict(table, html=_restore_image_refs(table["html"], pair)) for table in tables |
| ] |
|
|
|
|
| def _restore_image_refs(table_html, pair): |
| """Restore document page IDs, then resolve crops from re-rendered pages exactly.""" |
| local_to_document = { |
| local: number for local, number in enumerate(pair.page_numbers, 1) |
| } |
| table_html = remap(table_html, local_to_document) |
|
|
| def restore(match): |
| reference = match[0] |
| if match[1] is None: |
| raise ValueError("Output crop lacks a page number.") |
| if int(match[1]) in pair.rendered_pages: |
| if reference not in pair.alias_to_original: |
| raise ValueError("Unmappable crop from a rendered page: " + reference) |
| return pair.alias_to_original[reference] |
| return reference |
|
|
| return IMAGE_REF.sub(restore, table_html) |
|
|
|
|
| def _apply_merge(state, pair, tables): |
| """Validate writeback before updating page contents or fragment provenance.""" |
| earlier, later = (state.pages[number - 1] for number in pair.page_numbers) |
| before = (earlier.markdown, later.markdown) |
| after = apply_pair(before, tables) |
| if any(TABLE.sub("", old) != TABLE.sub("", new) for old, new in zip(before, after)): |
| raise ValueError("Writeback changed non-table content.") |
| if tables: |
| sources = earlier.table_sources[-1] | later.table_sources[0] |
| state.merged_tables = [ |
| group |
| for group in state.merged_tables |
| if not sources.intersection(map(tuple, group["fragment_ids"])) |
| ] |
| state.merged_tables.append( |
| { |
| "source_pages": sorted({number for number, _ in sources}), |
| "fragment_ids": sorted(sources), |
| "html": tables[0]["html"], |
| } |
| ) |
| earlier.table_sources[-1] = sources |
| del later.table_sources[0] |
| earlier.markdown, later.markdown = after |
|
|
|
|
| def _save_document_crops(markdown, images, output): |
| """Materialize page-qualified references from the original input images.""" |
| from PIL import Image |
|
|
| errors = [] |
| references = {} |
| for match in IMAGE_REF.finditer(markdown): |
| references.setdefault(match[1], {})[match[0]] = match.groups()[1:] |
| for page, crops in references.items(): |
| try: |
| if page is None or not 1 <= int(page) <= len(images): |
| raise ValueError("Image reference needs a valid source page number") |
| with Image.open(images[int(page) - 1]) as source: |
| for reference, coordinates in crops.items(): |
| try: |
| left, top, right, bottom = map(int, coordinates) |
| width, height = source.size |
| box = ( |
| max(0, min(width, round(left * width / 1000))), |
| max(0, min(height, round(top * height / 1000))), |
| max(0, min(width, round(right * width / 1000))), |
| max(0, min(height, round(bottom * height / 1000))), |
| ) |
| if box[2] <= box[0] or box[3] <= box[1]: |
| raise ValueError("Image crop has zero or negative pixel area") |
| target = output / reference |
| target.parent.mkdir(parents=True, exist_ok=True) |
| source.crop(box).convert("RGB").save(target) |
| except Exception as exc: |
| errors.append({"reference": reference, "message": str(exc)}) |
| except Exception as exc: |
| errors.extend( |
| {"reference": reference, "message": str(exc)} for reference in crops |
| ) |
| return errors |
|
|
|
|
| def _finalize_document(state, images, output, error): |
| document = "\n\n".join(page.markdown for page in state.pages) + "\n" |
| crop_errors = _save_document_crops(document, images, output) |
| for page in state.pages: |
| target = output / "pages" / f"page_{page.number:04d}.md" |
| target.parent.mkdir(exist_ok=True) |
| |
| target.write_text(IMAGE_REF.sub(lambda m: "../" + m[0], page.markdown)) |
| (output / "document.md").write_text(document) |
| if crop_errors and error is None: |
| error = {"type": "ImageCropExportFailed", "recoverable": False} |
| save(output / "failure.json", {**error, "image_export_errors": crop_errors}) |
| if error is None and state.parse_warnings: |
| error = { |
| "type": "Stage1OutputIncomplete", |
| "recoverable": True, |
| "pages": [warning["page"] for warning in state.parse_warnings], |
| } |
| result = { |
| "status": "COMPLETE" if error is None else "PARTIAL", |
| "error": error, |
| "input_pages": len(images), |
| "parsed_pages": len(state.pages), |
| "chunks": state.chunks, |
| "steps": state.steps, |
| "merged_tables": state.merged_tables, |
| "stage1_parse_warnings": state.parse_warnings, |
| "image_export_errors": crop_errors, |
| } |
| save(output / "result.json", result) |
| return result |
|
|
|
|
| |
|
|
|
|
| def _document_images(image_paths): |
| images = [str(Path(path).resolve()) for path in image_paths] |
| if not images or not all(Path(path).is_file() for path in images): |
| raise ValueError("Provide existing page images in reading order.") |
| return images |
|
|
|
|
| def save(path, value): |
| """Write a JSON artifact, creating its parent directory when needed.""" |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps(value, ensure_ascii=False, indent=2) + "\n") |
|
|
|
|
| def _is_valid_bbox(value): |
| if not isinstance(value, list) or len(value) != 4: |
| return False |
| if not all(type(coordinate) is int for coordinate in value): |
| return False |
| left, top, right, bottom = value |
| return 0 <= left <= right < 1000 and 0 <= top <= bottom < 1000 |
|
|
|
|
| def _recover_json_prefix(text): |
| """Return only complete prefix objects; never guess or repair JSON syntax.""" |
| blocks, unparsed_tail_chars = [], 0 |
| if not text.startswith("["): |
| return blocks, unparsed_tail_chars |
| position = 1 |
| decoder = json.JSONDecoder() |
| while position < len(text): |
| while position < len(text) and text[position].isspace(): |
| position += 1 |
| if position == len(text) or text[position] == "]": |
| break |
| try: |
| block, end = decoder.raw_decode(text, position) |
| except json.JSONDecodeError: |
| unparsed_tail_chars = len(text) - position |
| break |
| blocks.append(block) |
| position = end |
| while position < len(text) and text[position].isspace(): |
| position += 1 |
| if position < len(text) and text[position] == ",": |
| position += 1 |
| elif position < len(text) and text[position] != "]": |
| unparsed_tail_chars = len(text) - position |
| break |
| return blocks, unparsed_tail_chars |
|
|
|
|
| def structured_to_markdown(text): |
| """Project content in emitted order, retaining diagnostics for incomplete JSON. |
| |
| Preserve every category and keep only complete prefix objects on JSON failure. |
| Invalid bbox/category fields are reported without discarding valid content. |
| """ |
| stripped = text.strip() |
| if stripped.startswith("```") and stripped.endswith("```"): |
| stripped = re.sub(r"^```(?:json)?\s*\n|\n```$", "", stripped, flags=re.I) |
| diagnostics = { |
| "json_complete": False, |
| "schema_issues": [], |
| "invalid_content_objects": [], |
| "unparsed_tail_chars": 0, |
| } |
| try: |
| blocks = json.loads(stripped) |
| if not isinstance(blocks, list): |
| raise ValueError("Expected one JSON array") |
| diagnostics["json_complete"] = True |
| except ValueError as exc: |
| diagnostics["parse_error"] = str(exc) |
| blocks, diagnostics["unparsed_tail_chars"] = _recover_json_prefix(stripped) |
|
|
| contents = [] |
| for index, block in enumerate(blocks): |
| if not isinstance(block, dict) or not isinstance(block.get("content"), str): |
| diagnostics["invalid_content_objects"].append(index) |
| continue |
| contents.append(block["content"]) |
| if not _is_valid_bbox(block.get("bbox")): |
| diagnostics["schema_issues"].append({"index": index, "field": "bbox"}) |
| if ( |
| not isinstance(block.get("category"), str) |
| or block["category"] not in CATEGORIES |
| ): |
| diagnostics["schema_issues"].append({"index": index, "field": "category"}) |
| diagnostics.update( |
| parsed_objects=len(blocks), |
| emitted_contents=len(contents), |
| policy="join content in emitted array order; preserve every category; retain complete prefix objects only", |
| ) |
| diagnostics["valid"] = ( |
| diagnostics["json_complete"] |
| and not diagnostics["schema_issues"] |
| and not diagnostics["invalid_content_objects"] |
| and bool(contents) |
| ) |
| return "\n\n".join(contents), blocks, diagnostics |
|
|
|
|
| def _clean_truncated_repeats( |
| text: str, |
| min_text_len: int = 8000, |
| max_period: int = 200, |
| min_period: int = 1, |
| min_repeat_chars: int = 100, |
| min_repeat_times: int = 5, |
| ) -> str: |
| n = len(text) |
| if n < min_text_len: |
| return text |
|
|
| max_period = min(max_period, n - 1) |
| for unit_len in range(min_period, max_period + 1): |
| if text[n - 1] != text[n - 1 - unit_len]: |
| continue |
|
|
| match_len = 1 |
| idx = n - 2 |
| while idx >= unit_len and text[idx] == text[idx - unit_len]: |
| match_len += 1 |
| idx -= 1 |
|
|
| total_len = match_len + unit_len |
| repeat_times = total_len // unit_len |
| tail_len = total_len % unit_len |
|
|
| if repeat_times >= min_repeat_times and total_len >= min_repeat_chars: |
| return text[: n - total_len + unit_len] + text[n - tail_len :] |
|
|
| return text |
|
|
|
|
| def parse_tables(raw: str) -> list[dict]: |
| """Return all positive objects or [] for the exact negative sentinel.""" |
| items = json.loads(raw) |
| if not isinstance(items, list) or not items: |
| raise ValueError("Expected a nonempty JSON array.") |
| for item in items: |
| if not isinstance(item, dict) or set(item) != {"source_pages", "html"}: |
| raise ValueError("Each object needs exactly source_pages and html.") |
| if not isinstance(item["source_pages"], list) or not isinstance( |
| item["html"], str |
| ): |
| raise ValueError("Invalid field types.") |
| if not item["source_pages"]: |
| if len(items) == 1 and item["html"] == NEGATIVE: |
| return [] |
| raise ValueError("Invalid or mixed negative sentinel.") |
| if item["source_pages"] != [1, 2] or any( |
| type(p) is not int for p in item["source_pages"] |
| ): |
| raise ValueError("A positive two-page result must use [1, 2].") |
| if not TABLE.fullmatch(item["html"].strip()): |
| raise ValueError("Expected complete HTML table output.") |
| return items |
|
|
|
|
| def apply_pair( |
| markdowns: tuple[str, str], |
| tables: list[dict], |
| bindings: list[tuple[int, int]] | None = None, |
| ) -> tuple[str, str]: |
| """Replace page 1's last table and remove page 2's first table by default. |
| |
| Multiple merged tables require explicit zero-based fragment bindings. |
| Nested or incomplete HTML tables are rejected; non-table source is retained. |
| """ |
| if len(markdowns) != 2: |
| raise ValueError("Need exactly two pages.") |
| spans = [list(TABLE.finditer(md)) for md in markdowns] |
| if bindings is None: |
| if not tables: |
| return tuple(markdowns) |
| if len(tables) != 1: |
| raise ValueError("Multiple merged tables require explicit source bindings.") |
| bindings = [(len(spans[0]) - 1, 0)] |
| if len(tables) != len(bindings): |
| raise ValueError("Need one explicit binding per merged table.") |
| for md, matches in zip(markdowns, spans): |
| if len(re.findall(r"<table\b", md, re.I)) != len(matches) or len( |
| re.findall(r"</table\s*>", md, re.I) |
| ) != len(matches): |
| raise ValueError( |
| "Nested or incomplete tables require a layout-aware editor." |
| ) |
| edits = [{}, {}] |
| for item, binding in zip(tables, bindings): |
| parse_tables(json.dumps([item])) |
| if len(re.findall(r"<table\b", item["html"], re.I)) != 1: |
| raise ValueError( |
| "Nested or multiple tables cannot be written into one fragment." |
| ) |
| if item["source_pages"] != [1, 2] or len(binding) != 2: |
| raise ValueError("Only positive two-page tables can be written back.") |
| for page, index in enumerate(binding): |
| if type(index) is not int or not 0 <= index < len(spans[page]): |
| raise ValueError("Source table index out of range.") |
| if index in edits[page]: |
| raise ValueError("Two results claim the same source fragment.") |
| edits[page][index] = item["html"] if page == 0 else "" |
| output = [] |
| for page, md in enumerate(markdowns): |
| for index in sorted(edits[page], reverse=True): |
| span = spans[page][index] |
| md = md[: span.start()] + edits[page][index] + md[span.end() :] |
| output.append(md) |
| return tuple(output) |
|
|
|
|
| def find_chunks(pages): |
| """Maximal runs of >=2 pages with predicted HTML tables; 1-based IDs.""" |
| chunks, run = [], [] |
| for page, markdown in enumerate(pages, 1): |
| if TABLE.search(markdown): |
| run.append(page) |
| else: |
| if len(run) >= 2: |
| chunks.append(run) |
| run = [] |
| if len(run) >= 2: |
| chunks.append(run) |
| return chunks |
|
|
|
|
| def remap(markdown: str, mapping: Mapping[int, int]) -> str: |
| def replace(match): |
| page = int(match[1]) |
| if page not in mapping: |
| raise ValueError(f"Image page {page} is outside the current pair.") |
| return f"page_{mapping[page]}/images/" |
|
|
| return REF.sub(replace, markdown) |
|
|
|
|
| def page_markdown(raw, page): |
| """Accept Markdown, optionally fenced; qualify page-local crop references.""" |
| markdown = re.sub( |
| r"^```(?:markdown|html)?\s*\n|\n```\s*$", "", raw.strip(), flags=re.I |
| ) |
| if not markdown: |
| raise ValueError("Empty Stage 1 Markdown.") |
| |
| return IMAGE_REF.sub( |
| lambda m: f"page_{page}/images/bbox_" + "_".join(m.groups()[1:]) + ".jpg", |
| markdown, |
| ) |
|
|
|
|
| |
|
|
|
|
| class VllmBackend: |
| """Keep engine, sampling and request settings aligned with README infer(). |
| |
| Benchmark-specific scheduling overrides do not belong in this backend. |
| """ |
|
|
| def __init__(self, model_path): |
| from vllm import LLM, SamplingParams |
|
|
| self.model = LLM( |
| model=str(Path(model_path).resolve()), |
| tensor_parallel_size=1, |
| gpu_memory_utilization=0.8, |
| gdn_prefill_backend="triton", |
| limit_mm_per_prompt={"image": 2}, |
| ) |
| self.tokenizer = self.model.get_tokenizer() |
| self.sampling = SamplingParams(max_tokens=16384, temperature=0.0) |
|
|
| def __call__(self, paths, task_prompt, folder, stage): |
| from PIL import Image |
|
|
| images = [] |
| for path in paths: |
| with Image.open(path) as image: |
| images.append(image.convert("RGB")) |
| content = [{"type": "image"} for _ in images] |
| content.append({"type": "text", "text": task_prompt}) |
| prompt = self.tokenizer.apply_chat_template( |
| [{"role": "user", "content": content}], |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| request = { |
| "prompt": prompt, |
| "multi_modal_data": {"image": images}, |
| "mm_processor_kwargs": { |
| "images_kwargs": { |
| "min_pixels": 448 * 448, |
| "max_pixels": 2880 * 2880, |
| } |
| }, |
| } |
| save( |
| folder / "input.json", |
| { |
| "stage": stage, |
| "images": list(paths), |
| "prompt": prompt, |
| "mm_processor_kwargs": request["mm_processor_kwargs"], |
| }, |
| ) |
| output = self.model.generate([request], self.sampling)[0].outputs[0] |
| return GenerationResult(output.text, output.finish_reason) |
|
|
|
|
| class _BrowserRenderer: |
| def __init__(self, mathjax): |
| self.mathjax = mathjax |
| self.playwright = self.browser = self.renderer = None |
|
|
| def render(self, *args): |
| if self.renderer is None: |
| from playwright.sync_api import sync_playwright |
|
|
| self.playwright = sync_playwright().start() |
| self.browser = self.playwright.chromium.launch(headless=True) |
| self.renderer = Renderer(self.browser, self.mathjax) |
| return self.renderer.render(*args) |
|
|
| def close(self): |
| if self.browser: |
| self.browser.close() |
| if self.playwright: |
| self.playwright.stop() |
|
|
|
|
| |
|
|
|
|
| CROP = re.compile(r"page_(\d+)/images/bbox_(\d+)_(\d+)_(\d+)_(\d+)\.jpg") |
| _CROP_ASSET = re.compile( |
| r"(?:page_(\d+)/)?images/(?:bbox_)?(\d+)_(\d+)_(\d+)_(\d+)\.(?:jpg|jpeg|png)", re.I |
| ) |
| _TAG = re.compile( |
| r"<!--.*?-->|</?[A-Za-z][\w:-]*(?:[^>\'\"]|\"[^\"]*\"|\'[^\']*\')*>", re.S |
| ) |
| _SRC = re.compile(r"""(?<![\w:-])src\s*=\s*(?:"([^"]*)"|'([^']*)'|([^\s>]+))""", re.I) |
| _ALLOWED = set( |
| "p div span br hr h1 h2 h3 h4 h5 h6 table thead tbody tfoot tr td th caption colgroup col ul ol li dl dt dd blockquote pre code b strong em i u s del sup sub a img figure figcaption section article main center font".split() |
| ) |
| _DROP = set("script style iframe object embed link meta base template".split()) |
| _CSS = """ |
| html,body{margin:0;padding:0;background:#fff;color:#111} |
| #capture{padding:28px;box-sizing:border-box;font-family:Arial,'Noto Sans CJK SC',sans-serif;font-size:18px;line-height:1.4} |
| table{border-collapse:collapse;width:100%;margin:12px 0}td,th{border:1px solid #666;padding:5px;vertical-align:top;overflow-wrap:anywhere} |
| img{max-width:100%;height:auto}h1{font-size:28px}h2{font-size:24px}h3{font-size:21px}p{margin:10px 0} |
| pre{white-space:pre-wrap;overflow-wrap:anywhere;font:inherit}code{font-family:monospace;white-space:pre-wrap;overflow-wrap:anywhere} |
| .math-display{display:block;text-align:center;margin:10px 0}.math-inline{display:inline-block;vertical-align:middle}.math-fallback{white-space:pre-wrap;overflow-wrap:anywhere} |
| mjx-container{max-width:none}svg{overflow:visible}th{font-weight:600} |
| """ |
|
|
|
|
| @dataclass(frozen=True) |
| class RenderLimits: |
| """Configurable engineering limits, not validated OCR accuracy thresholds.""" |
|
|
| max_side: int = 32768 |
| max_render_pixels: int = 64_000_000 |
| max_source_chars: int = 4_000_000 |
| min_effective_font_px: float = 12.0 |
| min_effective_cell_height_px: float = 8.0 |
| timeout_ms: int = 30000 |
|
|
| def __post_init__(self): |
| for name, value in asdict(self).items(): |
| if not math.isfinite(value) or value <= 0: |
| raise ValueError(name + " must be finite and positive") |
|
|
|
|
| class RenderError(ValueError): |
| """Concise caller-facing error; details remain in the render.json sidecar.""" |
|
|
| def __init__(self, code, message): |
| self.code = code |
| super().__init__(message) |
|
|
|
|
| def _escaped(text, pos): |
| start = pos |
| while start and text[start - 1] == "\\": |
| start -= 1 |
| return (pos - start) % 2 == 1 |
|
|
|
|
| def _dollar_is_math(body, mode): |
| if mode == "all": |
| return True |
| if mode == "off" or not body or body != body.strip(): |
| return False |
| if re.search(r"\\[A-Za-z]+", body): |
| return True |
| |
| if re.fullmatch(r"[A-Za-z]", body): |
| return True |
| if re.search(r"[^\x00-\x7f]", body): |
| return bool( |
| re.fullmatch(r"[\w\s\d.,_^=+*/<>{}()−-]+", body) |
| and re.search(r"[_^=+*/<>−-]", body) |
| ) |
| if re.fullmatch(r"\d+(?:\.\d+)?\\[,;!: ](?:[A-Za-z]+)", body): |
| return True |
| words = re.findall(r"[A-Za-z]+", body) |
| if any( |
| len(w) > 1 |
| and w not in {"sin", "cos", "tan", "log", "ln", "exp", "max", "min", "lim"} |
| for w in words |
| ): |
| return False |
| return bool(re.search(r"[_^=+*/<>{}]", body) or re.search(r"\w\s*-\s*\w", body)) |
|
|
|
|
| def _markdown_image(text, pos, parser, env): |
| from markdown_it.common.utils import normalizeReference |
| from markdown_it.helpers import parseLinkDestination, parseLinkLabel, parseLinkTitle |
| from markdown_it.rules_inline.state_inline import StateInline |
|
|
| state = StateInline(text, parser, env, []) |
| end_label = parseLinkLabel(state, pos + 1, False) |
| if end_label < 0: |
| return None |
| alt = text[pos + 2 : end_label] |
| cursor = end_label + 1 |
| title = "" |
| if cursor < len(text) and text[cursor] == "(": |
| cursor += 1 |
| while cursor < len(text) and text[cursor].isspace(): |
| cursor += 1 |
| dest = parseLinkDestination(text, cursor, len(text)) |
| if not dest.ok: |
| return None |
| ref, cursor = dest.str, dest.pos |
| spaces = cursor |
| while cursor < len(text) and text[cursor].isspace(): |
| cursor += 1 |
| if cursor > spaces: |
| parsed_title = parseLinkTitle(text, cursor, len(text)) |
| if parsed_title.ok: |
| title, cursor = parsed_title.str, parsed_title.pos |
| while cursor < len(text) and text[cursor].isspace(): |
| cursor += 1 |
| if cursor >= len(text) or text[cursor] != ")": |
| return None |
| end = cursor + 1 |
| else: |
| label, end = alt, cursor |
| if cursor < len(text) and text[cursor] == "[": |
| close = text.find("]", cursor + 1) |
| if close < 0: |
| return None |
| label = text[cursor + 1 : close] or alt |
| end = close + 1 |
| definition = env.get("references", {}).get(normalizeReference(label)) |
| if not definition: |
| return None |
| ref, title = definition["href"], definition.get("title", "") |
| return end, ref, alt, title |
|
|
|
|
| def _image_markup(item, ref): |
| if item["kind"] == "html": |
| return html.escape(ref, quote=True) |
| title = ( |
| (' "' + item["title"].replace("\\", "\\\\").replace('"', '\\"') + '"') |
| if item["title"] |
| else "" |
| ) |
| return "![" + item["alt"] + "](<" + ref + ">" + title + ")" |
|
|
|
|
| def _edit_source(text, edits): |
| out, previous = [], 0 |
| for start, end, value in sorted(edits): |
| if start < previous: |
| raise RenderError( |
| "overlapping_source_spans", |
| "Cannot safely map overlapping source elements.", |
| ) |
| out.extend([text[previous:start], value]) |
| previous = end |
| out.append(text[previous:]) |
| return "".join(out) |
|
|
|
|
| def prepare_source(source, parser, dollar_math="auto", html_document=False): |
| """Protect math before CommonMark and identify each actual image occurrence. |
| |
| Code fences, indented code, inline code, HTML attributes and comments remain |
| literal. Markdown image references are expanded per occurrence, not per URL. |
| """ |
| from markdown_it.helpers import parseLinkDestination, parseLinkTitle |
|
|
| prefix = "RDR" + uuid.uuid4().hex.upper() |
| env = {} |
| tokens = [] if html_document else parser.parse(source, env) |
| offsets = [0] |
| for line in source.splitlines(keepends=True): |
| offsets.append(offsets[-1] + len(line)) |
| literal = { |
| offsets[t.map[0]]: offsets[min(t.map[1], len(offsets) - 1)] |
| for t in tokens |
| if t.type in {"fence", "code_block"} and t.map |
| } |
| for definition in env.get("references", {}).values(): |
| if definition.get("map"): |
| a, b = definition["map"] |
| literal[offsets[a]] = offsets[min(b, len(offsets) - 1)] |
| formulas, images, edits, notes = [], [], [], [] |
| i = 0 |
| while i < len(source): |
| if i in literal: |
| i = literal[i] |
| continue |
| if source[i] == "`" and not _escaped(source, i): |
| run = re.match(r"`+", source[i:])[0] |
| close = re.search( |
| r"(?<!`)" + re.escape(run) + r"(?!`)", source[i + len(run) :] |
| ) |
| if close: |
| i += len(run) + close.end() |
| continue |
| if source.startswith("](", i): |
| |
| dest = parseLinkDestination(source, i + 2, len(source)) |
| if dest.ok: |
| cursor = dest.pos |
| while cursor < len(source) and source[cursor].isspace(): |
| cursor += 1 |
| title = parseLinkTitle(source, cursor, len(source)) |
| if title.ok: |
| cursor = title.pos |
| while cursor < len(source) and source[cursor].isspace(): |
| cursor += 1 |
| if cursor < len(source) and source[cursor] == ")": |
| i = cursor + 1 |
| continue |
| if source[i] == "<": |
| tag = _TAG.match(source, i) |
| if tag: |
| raw = tag[0] |
| kind = re.match(r"<([A-Za-z][\w:-]*)", raw) |
| kind = kind[1].lower() if kind else "" |
| if kind in _DROP | {"code"}: |
| close = re.search( |
| r"</" + re.escape(kind) + r"\s*>", source[tag.end() :], re.I |
| ) |
| if close: |
| i = tag.end() + close.end() |
| continue |
| if kind == "img": |
| src = _SRC.search(raw) |
| if src: |
| group = next(j for j in (1, 2, 3) if src[j] is not None) |
| start, end = i + src.start(group), i + src.end(group) |
| item = dict( |
| start=start, |
| end=end, |
| ref=html.unescape(src[group]), |
| kind="html", |
| token=prefix + "IMG" + str(len(images)), |
| ) |
| images.append(item) |
| edits.append((start, end, item["token"])) |
| i = tag.end() |
| continue |
| if source.startswith("![", i) and not _escaped(source, i): |
| found = _markdown_image(source, i, parser, env) |
| if found: |
| end, ref, alt, title = found |
| item = dict( |
| start=i, |
| end=end, |
| ref=ref, |
| alt=alt, |
| title=title, |
| kind="markdown", |
| token=prefix + "IMG" + str(len(images)), |
| ) |
| images.append(item) |
| edits.append((i, end, _image_markup(item, item["token"]))) |
| i = end |
| continue |
| opening = next( |
| (s for s in ("\\(", "\\[", "$$", "$") if source.startswith(s, i)), None |
| ) |
| if opening and not _escaped(source, i): |
| closing = {"\\(": "\\)", "\\[": "\\]", "$$": "$$", "$": "$"}[opening] |
| pos = i + len(opening) |
| while True: |
| end = source.find(closing, pos) |
| if end < 0: |
| break |
| if not _escaped(source, end): |
| break |
| pos = end + len(closing) |
| |
| |
| boundary = re.search( |
| r"</?(?:table|thead|tbody|tfoot|tr|td|th|div|p|h[1-6]|li|ul|ol|pre|blockquote)\b[^>]*>|<br\b[^>]*>", |
| source[pos:], |
| re.I, |
| ) |
| if boundary and end >= 0 and pos + boundary.start() < end: |
| notes.append({"code": "formula_block_boundary", "offset": i}) |
| end = -1 |
| if end >= 0: |
| body = source[i + len(opening) : end] |
| allowed = True |
| if opening == "$": |
| allowed = ( |
| "\n" not in body |
| and not source.startswith("$$", end) |
| and _dollar_is_math(body, dollar_math) |
| ) |
| if not allowed: |
| notes.append( |
| {"code": "literal_or_ambiguous_dollar", "offset": i} |
| ) |
| elif "\n\n" in body and opening == "\\(": |
| allowed = False |
| if allowed: |
| token = prefix + "MATH" + str(len(formulas)) + "X" |
| formulas.append( |
| dict( |
| token=token, |
| tex=html.unescape(body), |
| raw=source[i : end + len(closing)], |
| display=opening in {"\\[", "$$"}, |
| ) |
| ) |
| edits.append((i, end + len(closing), token)) |
| i = end + len(closing) |
| continue |
| elif opening != "$": |
| |
| stop = source.find("\n\n", i) |
| stop = len(source) if stop < 0 else stop |
| tag_boundary = _TAG.search(source, i + len(opening)) |
| if tag_boundary: |
| stop = min(stop, tag_boundary.start()) |
| token = prefix + "MATH" + str(len(formulas)) + "X" |
| formulas.append( |
| dict( |
| token=token, |
| tex="", |
| raw=source[i:stop], |
| display=False, |
| malformed=True, |
| ) |
| ) |
| edits.append((i, stop, token)) |
| i = stop |
| continue |
| |
| i += len(opening) |
| continue |
| i += 1 |
| return _edit_source(source, edits), formulas, images, notes |
|
|
|
|
| def sanitize(soup): |
| from bs4 import Comment |
|
|
| removed = 0 |
| for node in list(soup.find_all(string=lambda t: isinstance(t, Comment))): |
| node.extract() |
| for tag in list(soup.find_all(True)): |
| if tag.name is None: |
| continue |
| if tag.name in {"svg", "canvas", "math", "video", "audio"}: |
| raise RenderError( |
| "unsupported_embedded_media", |
| "Embedded media needs an explicit image asset; source retained.", |
| ) |
| if tag.name in _DROP: |
| tag.decompose() |
| removed += 1 |
| continue |
| if tag.name not in _ALLOWED: |
| tag.unwrap() |
| removed += 1 |
| continue |
| attrs = {} |
| if tag.name in {"td", "th", "col", "colgroup"}: |
| for key in ("rowspan", "colspan", "span"): |
| if key in tag.attrs: |
| value = str(tag[key]) |
| if not value.isdecimal() or not ( |
| 0 if key == "rowspan" else 1 |
| ) <= int(value) <= (65534 if key == "rowspan" else 1000): |
| raise RenderError( |
| "invalid_table_span", |
| "Invalid table cell span; source retained.", |
| ) |
| attrs[key] = value |
| if tag.name == "img": |
| attrs.update({k: tag[k] for k in ("src", "alt", "title") if k in tag.attrs}) |
| if tag.name == "ol" and str(tag.get("start", "")).lstrip("-").isdigit(): |
| attrs["start"] = tag["start"] |
| |
| declarations = dict( |
| (k.strip().lower(), v.strip().lower()) |
| for k, v in re.findall(r"([\w-]+)\s*:\s*([^;]+)", str(tag.get("style", ""))) |
| ) |
| align = declarations.get("text-align", str(tag.get("align", "")).lower()) |
| valign = declarations.get("vertical-align", str(tag.get("valign", "")).lower()) |
| style = [] |
| if align in {"left", "right", "center", "justify"}: |
| style.append("text-align:" + align) |
| if valign in {"top", "middle", "bottom", "baseline"}: |
| style.append("vertical-align:" + valign) |
| if style: |
| attrs["style"] = ";".join(style) |
| if dict(tag.attrs) != attrs: |
| removed += 1 |
| tag.attrs = attrs |
| return removed |
|
|
|
|
| def _png_data(image): |
| from PIL import Image |
|
|
| image = image.convert("RGBA") |
| background = Image.new("RGBA", image.size, "white") |
| background.alpha_composite(image) |
| buf = io.BytesIO() |
| background.convert("RGB").save(buf, format="PNG") |
| return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("ascii") |
|
|
|
|
| def resolve_alias(ref, aliases, *, tolerance=2, min_iou=0.98, events=None): |
| """Exact first; a small coordinate jitter must match one unique occurrence. |
| |
| This is a configurable rounding tolerance, not semantic figure localization. |
| An ambiguous or newly predicted region is never assigned to the nearest image. |
| """ |
| if ref in aliases: |
| return aliases[ref] |
| target = CROP.fullmatch(ref) |
| if not target: |
| raise RenderError("unmapped_crop", "Cannot map an image region to its source.") |
| page, *box = map(int, target.groups()) |
| candidates = [] |
| for key in aliases: |
| m = CROP.fullmatch(key) |
| if not m or int(m[1]) != page: |
| continue |
| other = list(map(int, m.groups()[1:])) |
| if max(abs(a - b) for a, b in zip(box, other)) > tolerance: |
| continue |
| intersection = max(0, min(box[2], other[2]) - max(box[0], other[0])) * max( |
| 0, min(box[3], other[3]) - max(box[1], other[1]) |
| ) |
| area = max(0, box[2] - box[0]) * max(0, box[3] - box[1]) |
| area2 = max(0, other[2] - other[0]) * max(0, other[3] - other[1]) |
| union = area + area2 - intersection |
| if union and intersection / union >= min_iou: |
| candidates.append(key) |
| if len(candidates) != 1: |
| raise RenderError( |
| "ambiguous_crop" if candidates else "unmapped_crop", |
| "Cannot uniquely map an image region to its source.", |
| ) |
| if events is not None: |
| events.append( |
| { |
| "code": "coordinate_rounding_match", |
| "predicted": ref, |
| "matched": candidates[0], |
| } |
| ) |
| return aliases[candidates[0]] |
|
|
|
|
| class Renderer: |
| def __init__( |
| self, |
| browser, |
| mathjax, |
| *, |
| limits=None, |
| dollar_math="auto", |
| assets=None, |
| asset_root=None, |
| resize=None, |
| resize_name=None, |
| max_input_pixels=None, |
| ): |
| from markdown_it import MarkdownIt |
|
|
| if dollar_math not in {"auto", "all", "off"}: |
| raise ValueError("dollar_math must be auto, all or off") |
| self.mathjax = Path(mathjax).resolve() |
| self.limits = limits or RenderLimits() |
| self.dollar_math = dollar_math |
| self.assets = dict(assets or {}) |
| self.asset_root = Path(asset_root).resolve() if asset_root else None |
| self.resize = resize |
| self.resize_name = resize_name or ( |
| "custom" |
| if resize |
| else "area_budget_estimate" if max_input_pixels else "unknown" |
| ) |
| if max_input_pixels is not None and max_input_pixels <= 0: |
| raise ValueError("max_input_pixels must be positive") |
| self.max_input_pixels = max_input_pixels |
| self.page = browser.new_page( |
| viewport={"width": 1200, "height": 900}, device_scale_factor=1 |
| ) |
| self.page.set_default_timeout(self.limits.timeout_ms) |
| |
| self.page.route("**/*", lambda route: route.abort()) |
| self.md = MarkdownIt("commonmark", {"html": True}).enable( |
| ["table", "strikethrough"] |
| ) |
|
|
| def _asset(self, ref, owner, paths): |
| from PIL import Image |
|
|
| crop = _CROP_ASSET.fullmatch(ref) |
| if crop: |
| page, left, top, right, bottom = crop.groups() |
| page = int(page) if page else owner |
| left, top, right, bottom = map(int, (left, top, right, bottom)) |
| if ( |
| not 1 <= page <= len(paths) |
| or not 0 <= left < right <= 1000 |
| or not 0 <= top < bottom <= 1000 |
| ): |
| raise RenderError( |
| "invalid_crop", |
| "An image crop has invalid coordinates; source retained.", |
| ) |
| with Image.open(paths[page - 1]) as im: |
| box = ( |
| round(left * im.width / 1000), |
| round(top * im.height / 1000), |
| round(right * im.width / 1000), |
| round(bottom * im.height / 1000), |
| ) |
| if box[0] >= box[2] or box[1] >= box[3]: |
| raise RenderError( |
| "empty_crop", "An image crop is smaller than one pixel." |
| ) |
| if (box[2] - box[0]) * ( |
| box[3] - box[1] |
| ) > self.limits.max_render_pixels: |
| raise RenderError( |
| "asset_limit", |
| "An image asset exceeds the rendering pixel budget.", |
| ) |
| return _png_data(im.crop(box)) |
| if ref.startswith("data:"): |
| header, sep, payload = ref.partition(",") |
| if not sep or not re.fullmatch( |
| r"data:image/(?:png|jpeg|webp|gif);base64", header, re.I |
| ): |
| raise RenderError( |
| "unsupported_asset", |
| "An image format is unsupported; source retained.", |
| ) |
| with Image.open(io.BytesIO(base64.b64decode(payload, validate=True))) as im: |
| if im.width * im.height > self.limits.max_render_pixels: |
| raise RenderError( |
| "asset_limit", |
| "An image asset exceeds the rendering pixel budget.", |
| ) |
| return _png_data(im) |
| path = self.assets.get(ref) |
| if path is None and self.asset_root: |
| from urllib.parse import unquote, urlsplit |
|
|
| if not urlsplit(ref).scheme and not ref.startswith("//"): |
| candidate = (self.asset_root / unquote(ref)).resolve() |
| if candidate.is_relative_to(self.asset_root): |
| path = candidate |
| if path is None: |
| raise RenderError( |
| "unresolved_asset", |
| "An image needs a local asset mapping; source retained.", |
| ) |
| with Image.open(path) as im: |
| if im.width * im.height > self.limits.max_render_pixels: |
| raise RenderError( |
| "asset_limit", "An image asset exceeds the rendering pixel budget." |
| ) |
| return _png_data(im) |
|
|
| def render(self, md, owner, paths, directory, name): |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| if Path(name).name != name or not name: |
| raise ValueError("name must be a filename stem") |
| started = time.monotonic() |
| diag = { |
| "version": 2, |
| "status": "ERROR", |
| "math": 0, |
| "errors": 0, |
| "broken": 0, |
| "images": [], |
| "notes": [], |
| "limits": asdict(self.limits), |
| "resize_policy": self.resize_name, |
| } |
| (directory / (name + ".source.md")).write_text(md) |
| try: |
| result = self._render(md, owner, paths, directory, name, diag) |
| diag["status"] = "DEGRADED" if diag.get("formula_fallbacks") else "COMPLETE" |
| return result |
| except Exception as exc: |
| diag["error"] = { |
| "code": getattr(exc, "code", "render_failed"), |
| "type": type(exc).__name__, |
| "message": str(exc)[:1000], |
| } |
| if isinstance(exc, RenderError): |
| raise |
| raise RenderError( |
| "render_failed", |
| "Page rendering failed; source and diagnostics were retained.", |
| ) from exc |
| finally: |
| diag["seconds"] = time.monotonic() - started |
| save(directory / (name + ".render.json"), diag) |
|
|
| def _render(self, md, owner, paths, directory, name, diag): |
| from bs4 import BeautifulSoup, Comment, NavigableString |
| from PIL import Image |
|
|
| if not 1 <= owner <= len(paths): |
| raise RenderError("invalid_owner", "Page index is outside the document.") |
| if len(md) > self.limits.max_source_chars: |
| raise RenderError( |
| "source_limit", "Page source exceeds the configured rendering limit." |
| ) |
| probe = ( |
| BeautifulSoup(md, "html.parser") if md.lstrip().startswith("<") else None |
| ) |
| html_document = bool( |
| probe is not None |
| and probe.find(True) |
| and all( |
| not isinstance(x, NavigableString) |
| or isinstance(x, Comment) |
| or not str(x).strip() |
| for x in probe.contents |
| ) |
| ) |
| prepared, formulas, occurrences, notes = prepare_source( |
| md, self.md, self.dollar_math, html_document |
| ) |
| diag["input_format"] = "html" if html_document else "markdown" |
| diag["notes"].extend(notes) |
| soup = BeautifulSoup( |
| prepared if html_document else self.md.render(prepared), "html.parser" |
| ) |
| diag["sanitized_elements"] = sanitize(soup) |
| by_token = {x["token"]: x for x in occurrences} |
| seen = set() |
| for image in soup.find_all("img"): |
| key = image.get("src", "") |
| if key not in by_token or key in seen: |
| raise RenderError( |
| "image_mapping", "Cannot safely map image occurrences." |
| ) |
| seen.add(key) |
| item = by_token[key] |
| image["src"] = self._asset(item["ref"], owner, paths) |
| image["data-image-id"] = key |
| if len(seen) != len(occurrences): |
| raise RenderError( |
| "image_mapping", "An image occurrence was lost during parsing." |
| ) |
| by_math = {x["token"]: x for x in formulas} |
| if by_math: |
| pattern = re.compile( |
| "|".join(map(re.escape, sorted(by_math, key=len, reverse=True))) |
| ) |
| found = [] |
| for node in list(soup.find_all(string=pattern)): |
| if isinstance(node, Comment): |
| continue |
| pieces = pattern.split(str(node)) |
| matches = list(pattern.finditer(str(node))) |
| for k, part in enumerate(pieces): |
| if part: |
| node.insert_before(NavigableString(part)) |
| if k < len(matches): |
| item = by_math[matches[k][0]] |
| found.append(item["token"]) |
| tag = soup.new_tag("span") |
| tag["data-math-id"] = item["token"] |
| tag["class"] = ( |
| "math-display" if item["display"] else "math-inline" |
| ) |
| tag.string = item["raw"] |
| node.insert_before(tag) |
| node.extract() |
| if sorted(found) != sorted(by_math): |
| raise RenderError( |
| "formula_mapping", "A formula was lost during parsing." |
| ) |
| tables = soup.find_all("table") |
| expected_cells = len(soup.find_all(["td", "th"])) |
| columns = max( |
| ( |
| sum( |
| int(c.get("colspan", 1)) |
| for c in row.find_all(["td", "th"], recursive=False) |
| ) |
| for row in soup.find_all("tr") |
| ), |
| default=1, |
| ) |
| with Image.open(paths[owner - 1]) as original: |
| width = max(original.width, columns * 120 + 56) |
| document = ( |
| '<!doctype html><html><head><meta charset="utf-8"><style>' |
| + _CSS |
| + '</style></head><body><div id="capture" style="width:' |
| + str(width) |
| + 'px">' |
| + str(soup) |
| + "</div></body></html>" |
| ) |
| |
| (directory / (name + ".html")).write_text(document) |
| if width > self.limits.max_side: |
| raise RenderError( |
| "size_limit", |
| "Rendered page exceeds the configured size limit; complete source retained.", |
| ) |
| page = self.page |
| page.set_viewport_size({"width": width, "height": 900}) |
| page.set_content(document, wait_until="domcontentloaded") |
| diag["expected_formulas"] = len(formulas) |
| formula_results = [] |
| if formulas: |
| if not self.mathjax.is_file(): |
| raise RenderError("mathjax_missing", "Local MathJax bundle is missing.") |
| page.evaluate( |
| "window.MathJax={startup:{typeset:false},svg:{fontCache:'none'},tex:{maxBuffer:100000,packages:['base','ams','newcommand']}}" |
| ) |
| page.add_script_tag(path=str(self.mathjax)) |
| page.wait_for_function( |
| "window.MathJax && MathJax.startup && MathJax.startup.promise" |
| ) |
| formula_results = page.evaluate( |
| """async ({items,timeout})=>{ |
| const work=async()=>{ |
| await MathJax.startup.promise; |
| if(typeof MathJax.tex2svgPromise!=='function') throw Error('MathJax SVG API unavailable'); |
| const results=[]; |
| for(const item of items){ |
| const target=document.querySelector('[data-math-id="'+item.token+'"]'); |
| let error=null; |
| try{ |
| if(item.malformed) throw Error('Unclosed formula delimiter'); |
| if(!item.tex.trim()) throw Error('Empty formula'); |
| const rendered=await MathJax.tex2svgPromise(item.tex,{display:item.display}); |
| if(rendered.querySelector('[data-mml-node="merror"],mjx-merror,[data-mjx-error]')) throw Error('Invalid TeX'); |
| // This screenshot renderer needs only SVG. Standalone tex2svg |
| // may emit assistive MathML without the CSS that hides it. |
| rendered.querySelectorAll('mjx-assistive-mml').forEach(node=>node.remove()); |
| target.replaceChildren(rendered); |
| }catch(e){error=String(e.message||e);target.textContent=item.raw;target.classList.add('math-fallback');} |
| results.push({id:item.token,rendered:!error,error}); |
| } |
| return results; |
| }; |
| let timer;try{return await Promise.race([work(),new Promise((_,reject)=>{timer=setTimeout(()=>reject(Error('Formula rendering timeout')),timeout)})]);}finally{clearTimeout(timer);} |
| }""", |
| {"items": formulas, "timeout": self.limits.timeout_ms}, |
| ) |
| diag["formulas"] = formula_results |
| diag["formula_fallbacks"] = sum(not x["rendered"] for x in formula_results) |
| page.evaluate( |
| """async timeout=>{ |
| let timer;try{await Promise.race([(async()=>{await document.fonts.ready;await Promise.all([...document.images].map(i=>i.decode()));await new Promise(r=>requestAnimationFrame(()=>requestAnimationFrame(r)));})(),new Promise((_,reject)=>{timer=setTimeout(()=>reject(Error('Image/font readiness timeout')),timeout)})]);}finally{clearTimeout(timer);} |
| }""", |
| self.limits.timeout_ms, |
| ) |
| |
| needed = page.evaluate("document.documentElement.scrollWidth") |
| if needed > width: |
| width = needed + 28 |
| if width > self.limits.max_side: |
| raise RenderError( |
| "size_limit", |
| "Rendered content is too wide; complete source retained.", |
| ) |
| page.set_viewport_size({"width": width, "height": 900}) |
| page.locator("#capture").evaluate('(e,w)=>e.style.width=w+"px"', width) |
| layout = page.evaluate("""()=>{ |
| const root=document.getElementById('capture');const area=root.getBoundingClientRect(); |
| const boxes=[...root.querySelectorAll('table,td,th,img,[data-math-id]')]; |
| const bad=boxes.filter(e=>{const b=e.getBoundingClientRect(),s=getComputedStyle(e);return !b.width||!b.height||s.display==='none'||s.visibility==='hidden'||Number(s.opacity)===0||b.left<0||b.right>area.right+.5;}).map(e=>({tag:e.tagName,id:e.dataset.imageId||e.dataset.mathId||null})); |
| const fonts=[];const walker=document.createTreeWalker(root,NodeFilter.SHOW_TEXT);let n; |
| while(n=walker.nextNode()){if(n.textContent.trim()&&!n.parentElement.closest('mjx-container')){const s=getComputedStyle(n.parentElement);fonts.push(parseFloat(s.fontSize));}} |
| const cellHeights=[...root.querySelectorAll('td,th')].map(e=>e.getBoundingClientRect().height); |
| return {math:root.querySelectorAll('mjx-container').length,broken:[...root.querySelectorAll('img')].filter(i=>!i.complete||!i.naturalWidth).length, |
| tables:root.querySelectorAll('table').length,cells:cellHeights.length,invisible_or_clipped:bad, |
| min_font_px:fonts.length?fonts.reduce((a,b)=>Math.min(a,b),Infinity):18,min_cell_height_px:cellHeights.length?cellHeights.reduce((a,b)=>Math.min(a,b),Infinity):null, |
| width:document.documentElement.scrollWidth,height:document.documentElement.scrollHeight, |
| images:[...root.querySelectorAll('img')].map(i=>{const b=i.getBoundingClientRect();return {id:i.dataset.imageId,x:b.x,y:b.y,w:b.width,h:b.height}})}; |
| }""") |
| diag.update(layout) |
| diag["errors"] = diag["formula_fallbacks"] |
| if layout["tables"] != len(tables) or layout["cells"] != expected_cells: |
| raise RenderError( |
| "table_structure", |
| "Browser parsing changed the table structure; source retained.", |
| ) |
| if layout["invisible_or_clipped"] or layout["broken"]: |
| raise RenderError( |
| "invisible_content", |
| "Some page content is not visible; source retained.", |
| ) |
| if layout["math"] + diag["formula_fallbacks"] != len(formulas): |
| raise RenderError( |
| "formula_coverage", "Formula rendering coverage is incomplete." |
| ) |
| w, h = layout["width"], layout["height"] |
| |
| (directory / (name + ".html")).write_text( |
| page.evaluate( |
| "() => { const root=document.documentElement.cloneNode(true); root.querySelectorAll('script').forEach(e=>e.remove()); return '<!doctype html>'+root.outerHTML; }" |
| ) |
| ) |
| if max(w, h) > self.limits.max_side or w * h > self.limits.max_render_pixels: |
| raise RenderError( |
| "size_limit", |
| "Rendered page exceeds the configured size limit; complete HTML retained.", |
| ) |
| target = (w, h) |
| if not self.resize and not self.max_input_pixels: |
| diag["notes"].append({"code": "model_resize_unspecified"}) |
| if self.resize: |
| try: |
| target = tuple(self.resize(w, h)) |
| except Exception as exc: |
| raise RenderError( |
| "model_resize_limit", |
| "Page exceeds the model image-size limits; complete HTML retained.", |
| ) from exc |
| elif self.max_input_pixels: |
| scale = min(1, math.sqrt(self.max_input_pixels / (w * h))) |
| target = (max(1, math.floor(w * scale)), max(1, math.floor(h * scale))) |
| if len(target) != 2 or min(target) <= 0: |
| raise RenderError( |
| "resize_policy", "Invalid target dimensions from resize policy." |
| ) |
| scale = min(target[0] / w, target[1] / h) |
| font = layout["min_font_px"] * scale |
| cell = ( |
| layout["min_cell_height_px"] * scale |
| if layout["min_cell_height_px"] |
| else None |
| ) |
| diag["readability"] = { |
| "input_size": [w, h], |
| "model_size": list(target), |
| "scale": scale, |
| "effective_font_px": font, |
| "effective_cell_height_px": cell, |
| "model_resize_known": bool(self.resize), |
| "estimated_from_pixel_budget": bool( |
| not self.resize and self.max_input_pixels |
| ), |
| "thresholds_are_heuristic": True, |
| } |
| if font < self.limits.min_effective_font_px or ( |
| cell is not None and cell < self.limits.min_effective_cell_height_px |
| ): |
| raise RenderError( |
| "readability_limit", |
| "Page text would be too small after resizing; complete HTML retained.", |
| ) |
| aliases, edits = {}, [] |
| for image in layout["images"]: |
| item = by_token[image["id"]] |
| image["ref"] = item["ref"] |
| box = [ |
| round(image["x"] / w * 1000), |
| round(image["y"] / h * 1000), |
| round((image["x"] + image["w"]) / w * 1000), |
| round((image["y"] + image["h"]) / h * 1000), |
| ] |
| if not 0 <= box[0] < box[2] <= 1000 or not 0 <= box[1] < box[3] <= 1000: |
| raise RenderError( |
| "crop_quantization", |
| "An image is too small for normalized coordinate mapping.", |
| ) |
| alias = f"page_{owner}/images/bbox_" + "_".join(map(str, box)) + ".jpg" |
| if alias in aliases: |
| raise RenderError( |
| "crop_alias_collision", |
| "Image positions cannot be uniquely represented.", |
| ) |
| aliases[alias] = item["ref"] |
| edits.append((item["start"], item["end"], _image_markup(item, alias))) |
| rewritten = _edit_source(md, edits) |
| path = directory / (name + ".png") |
| page.screenshot(path=str(path), full_page=True, timeout=self.limits.timeout_ms) |
| with Image.open(path) as rendered: |
| if rendered.size != (w, h): |
| raise RenderError( |
| "screenshot_dimensions", |
| "Screenshot dimensions differ from image coordinate mapping.", |
| ) |
| (directory / (name + ".rewritten.md")).write_text(rewritten) |
| save(directory / (name + ".aliases.json"), aliases) |
| return RenderResult(str(path), rewritten, aliases, diag) |
|
|
|
|
| |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser( |
| description="OvisOCR3 document-level cross-page reconstruction" |
| ) |
| parser.add_argument("images", nargs="+", help="Page images in reading order") |
| parser.add_argument("--model", default=".") |
| parser.add_argument("--output", required=True) |
| parser.add_argument( |
| "--stage1-prompt", |
| default="document_structured.txt", |
| choices=[ |
| "document_markdown", |
| "document_markdown.txt", |
| "document_structured", |
| "document_structured.txt", |
| ], |
| ) |
| parser.add_argument( |
| "--mathjax", |
| default=None, |
| help="Local MathJax bundle; defaults to assets/mathjax.js beside this script", |
| ) |
| args = parser.parse_args() |
| result = reconstruct_document( |
| args.images, |
| args.output, |
| model_path=args.model, |
| stage1_prompt=args.stage1_prompt, |
| mathjax=args.mathjax, |
| ) |
| print(json.dumps(result, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
