| |
| |
| """Prepare public ParseBench PDFs and annotations as compact evaluation shards.""" |
| from __future__ import annotations |
|
|
| import argparse |
| from concurrent.futures import ProcessPoolExecutor |
| import hashlib |
| import io |
| import json |
| import multiprocessing |
| from pathlib import Path |
| import tarfile |
|
|
| FILES = ['chart.jsonl', 'layout.jsonl', 'table.jsonl', 'text_content.jsonl', 'text_formatting.jsonl'] |
| SNAPSHOT = None |
|
|
|
|
| def load_groups(snapshot): |
| groups = {} |
| for name in FILES: |
| with (snapshot / name).open() as stream: |
| for line in stream: |
| if not line.strip(): |
| continue |
| row = json.loads(line) |
| relative = row['pdf'] |
| group = groups.setdefault(relative, {'categories': set(), 'rules': [], 'ground_truth': ''}) |
| category = row.get('category') or Path(name).stem |
| group['categories'].add(category) |
| details = row.get('rule') or {} |
| if isinstance(details, str): |
| details = json.loads(details) |
| rule = {'type': row['type'], **details} |
| for field in ['id', 'page', 'tags', 'category']: |
| if row.get(field) is not None and row.get(field) != []: |
| rule[field] = row[field] |
| group['rules'].append(rule) |
| if row.get('expected_markdown') and not group['ground_truth']: |
| group['ground_truth'] = row['expected_markdown'] |
| return sorted(groups.items()) |
|
|
|
|
| def render(task): |
| import pymupdf as fitz |
| index, relative = task |
| relative_path = Path(relative) |
| if relative_path.is_absolute() or '..' in relative_path.parts: |
| raise ValueError('Source outside dataset directory: ' + relative) |
| |
| source = SNAPSHOT / relative_path |
| with fitz.open(source) as doc: |
| if len(doc) < 1: |
| raise ValueError('Empty document: ' + relative) |
| pix = doc[0].get_pixmap(dpi=150, alpha=False) |
| image = pix.tobytes('jpeg', jpg_quality=92) |
| return index, relative, image, pix.width, pix.height, len(doc) |
|
|
|
|
| def add(archive, name, data): |
| info = tarfile.TarInfo(name) |
| info.size = len(data) |
| info.mtime = 0 |
| info.mode = 0o644 |
| archive.addfile(info, io.BytesIO(data)) |
|
|
|
|
| def main(): |
| global SNAPSHOT |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--snapshot', type=Path, required=True) |
| parser.add_argument('--output', type=Path, required=True) |
| parser.add_argument('--workers', type=int, default=8) |
| args = parser.parse_args() |
| SNAPSHOT = args.snapshot |
| groups = load_groups(SNAPSHOT) |
| if len(groups) != 2078: |
| raise ValueError(f'Expected 2078 unique pages, got {len(groups)}') |
| import pymupdf as fitz |
| if fitz.VersionBind != '1.26.3': |
| raise SystemExit('Use PyMuPDF==1.26.3 for the recorded rendering contract') |
| train = args.output / 'train' |
| if train.exists() and any(train.iterdir()): |
| raise SystemExit('Choose a new output directory; existing shards are preserved') |
| train.mkdir(parents=True, exist_ok=True) |
| manifest = [] |
| archive = None |
| try: |
| with ProcessPoolExecutor(max_workers=args.workers, mp_context=multiprocessing.get_context('fork')) as pool: |
| for index, relative, image, width, height, source_pages in pool.map(render, enumerate(name for name, group in groups), chunksize=4): |
| if index % 500 == 0: |
| if archive is not None: |
| archive.close() |
| archive = tarfile.open(train / f'part-{index // 500:06d}.tar', 'w') |
| key = f'parsebench_{index:05d}' |
| group = groups[index][1] |
| category = sorted(group['categories'])[0] |
| expected = group['ground_truth'] |
| annotation = {'block_boxes': [[0, 0, width, height]], 'block_classes': ['Text'], |
| 'block_text': [expected], 'ground_truth': expected, 'rules': group['rules'], |
| 'category': category, 'page_info': {'width': width, 'height': height}} |
| add(archive, key + '.jpg', image) |
| add(archive, key + '.json', json.dumps(annotation, ensure_ascii=False).encode()) |
| manifest.append({'key': key, 'source_pdf': relative, 'image_name': key + '.jpg', |
| 'width': width, 'height': height, 'source_pages': source_pages, 'source_page_index': 0, |
| 'image_sha256': hashlib.sha256(image).hexdigest()}) |
| if (index + 1) % 200 == 0: |
| print('Prepared', index + 1, '/', len(groups), flush=True) |
| finally: |
| if archive is not None: |
| archive.close() |
| (args.output / 'source_manifest.jsonl').write_text(''.join(json.dumps(row) + '\n' for row in manifest)) |
| receipt = {'complete': True, 'pages': len(manifest), 'renderer': 'PyMuPDF', 'renderer_version': fitz.VersionBind, |
| 'dpi': 150, 'jpeg_quality': 92, 'source_page_index': 0, |
| 'multi_page_sources': [row['source_pdf'] for row in manifest if row['source_pages'] > 1], |
| 'annotations_used_only_for_dataset_preparation_and_scoring': True, |
| 'annotation_files': {name: hashlib.sha256((SNAPSHOT / name).read_bytes()).hexdigest() for name in FILES}, |
| 'shards': {path.name: {'bytes': path.stat().st_size, 'sha256': hashlib.sha256(path.read_bytes()).hexdigest()} for path in sorted(train.glob('*.tar'))}} |
| (args.output / 'preparation.json').write_text(json.dumps(receipt, indent=2) + '\n') |
| print(json.dumps(receipt), flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
