| """Convert HPD-Parsing predictions (JSON) into per-page markdown for OmniDocBench. |
| |
| Input : JSON, a list of ``{img_path, pred}`` (pred is the ``<BLOCK> <type> [bbox] |
| <CHILD> <content>`` stream from ``document parsing with fork.``). |
| Output: a folder of ``<image_stem>.md`` files matching the OmniDocBench GT paths. |
| |
| python hpd_to_markdown.py --input preds.json --out-md pred_md/ \ |
| --simplify-left-paren --clean-formula-tail --norm-formula-flag --wrap-cjk-arith |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import re |
| from pathlib import Path |
|
|
|
|
| _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 remove_block_fork_tags(result, simplify_left_paren=True, clean_formula_tail_flag=True, |
| norm_formula_flag=True): |
| """Split on `<BLOCK>`, keep the text after each `<CHILD>`, and join in reading order.""" |
| seg_pattern = re.compile(r'[^<]*<CHILD>(.*)', re.DOTALL) |
| lines = [] |
| for seg in result.split('<BLOCK>')[1:]: |
| cat_m = re.match(r'\s*([a-zA-Z_]+)', seg) |
| if cat_m and cat_m.group(1).lower() in ['chart', 'seal']: |
| continue |
| m = seg_pattern.match(seg) |
| if not m: |
| continue |
| text = m.group(1).strip() |
| text = re.sub(r'\b\w+\s*\[\s*[-\d.,\s]+\]\s*<(?:FORK|CHILD|BLOCK)>', '', text) |
| text = re.sub(r'<(?:FORK|CHILD|BLOCK)>', '', 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]': |
| continue |
| 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) |
| lines.append(text) |
| return '\n\n'.join(lines).strip() |
|
|
|
|
| def basename_to_md_name(img_path: str) -> str: |
| return os.path.splitext(os.path.basename(img_path))[0] + ".md" |
|
|
|
|
| def convert_json(in_path, out_md_dir, simplify_left_paren=True, |
| clean_formula_tail_flag=True, norm_formula_flag=True) -> int: |
| with open(in_path, "r", encoding="utf-8") as f: |
| rows = json.load(f) |
| os.makedirs(out_md_dir, exist_ok=True) |
| n = 0 |
| for row in rows: |
| img_path = row.get("img_path") or row.get("image_path") |
| pred = row.get("pred") or row.get("prediction") or "" |
| if not img_path: |
| continue |
| md = remove_block_fork_tags(pred, simplify_left_paren, clean_formula_tail_flag, norm_formula_flag) |
| with open(os.path.join(out_md_dir, basename_to_md_name(img_path)), "w", encoding="utf-8") as f: |
| f.write(md) |
| n += 1 |
| return n |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--input", required=True, help="json path (list of {img_path, pred})") |
| ap.add_argument("--out-md", required=True, help="output markdown folder") |
| args = ap.parse_args() |
|
|
| if Path(args.input).suffix.lower() != ".json": |
| raise SystemExit(f"unsupported extension: {Path(args.input).suffix} (expects .json)") |
| n = convert_json(args.input, args.out_md) |
| print(f"[ok] wrote {n} markdown files -> {args.out_md}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
