File size: 3,512 Bytes
dca5b2d e7ea792 dca5b2d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Prediction-only merge and Markdown conversion; no labels or scorer imports."""
from __future__ import annotations
import argparse,gzip,json,re,time,multiprocessing
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from types import SimpleNamespace
from engine import merge,markdown as adapter,policies
from engine.header_inference_probe import infer_header_rows
from engine.text_cleanup import normalize_text,repair_empty_prefix,dedupe_exact_elements
from engine.caption_relocation import relocate_above_captions
from engine.chart_context import add_native_subtitles
from engine.chart_caption import transform
def read_rows(path):
with (gzip.open(path,'rt') if str(path).endswith('.gz') else open(path)) as f:
return [json.loads(s) for s in f if s.strip()]
def write_rows(path,rows):
path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
with (gzip.open(path,'wt') if str(path).endswith('.gz') else path.open('w')) as f:
for row in rows:f.write(json.dumps(row,ensure_ascii=False)+'\n')
def key(row):return Path(row.get('key') or row.get('image_name','')).stem.removeprefix('parsebench__')
def arguments(crops=None):
return SimpleNamespace(crop_element_predictions=[Path(crops)] if crops else [],structural_classes='Table,Chart',
include_detector_labels='chart,table',max_crop_output_chars=0,max_crop_output_blocks=0,min_structural_area=0.,
crop_coordinate_frame='content_normalized',crop_pad_position='top-left',prediction_pad_position='top-left',
output_coordinate_frame='content_normalized',canvas_width=1664.,canvas_height=2048.,structural_source='crop',
dedupe_overlap_threshold=.85,dedupe_text_threshold=.8,replacement_text_threshold=.8)
def corrected(raw):
text=adapter.mmd_to_markdown(raw,text_transform=normalize_text)
return re.sub(r'<table\b[^>]*>.*?</table>',lambda m:infer_header_rows(m[0])[0],text,flags=re.S|re.I)
def _markdown_base(raw):
raw=policies.chart_render(raw,panels=True,compound=True)
def title(m):
if m[4]!='Chart':return m[0]
html,n=transform(corrected(m[0]))
return m[1]+html+m[3] if n else m[0]
return policies.promote_bold_header_row(corrected(policies.BLOCK.sub(title,raw)))
def markdown(raw):
text=_markdown_base(raw)
return add_native_subtitles(relocate_above_captions(raw,text,_markdown_base))
def finalize_one(task):
row,candidates=task;k=key(row)
raw=policies.replace_guarded(row['raw_out'],candidates,merge,guard='columns',repair=True,emit_repair=False)
raw=dedupe_exact_elements(repair_empty_prefix(raw))
return {'key':k,'raw_out':raw,'markdown':markdown(raw),'coordinate_frame':'content_normalized'}
def process(direct,crop_path=None,workers=8):
from parse_results import process as parse_process
return parse_process(direct,crop_path,workers)
def main():
p=argparse.ArgumentParser(description=__doc__);p.add_argument('--direct',type=Path,required=True)
p.add_argument('--crops',type=Path);p.add_argument('--output',type=Path,required=True);p.add_argument('--workers',type=int,default=8);a=p.parse_args()
start=time.monotonic();rows=process(read_rows(a.direct),a.crops,a.workers);write_rows(a.output,rows)
print(json.dumps({'pages':len(rows),'seconds':time.monotonic()-start,'output':str(a.output)}))
if __name__=='__main__':main()
|