Spaces:
Paused
Paused
File size: 4,758 Bytes
f66643d | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Group OmniDocBench page images into multi-page PDFs (default 32 images/PDF).
Why: we want to benchmark the *real* ``StageAParser.parse_pdf`` path. parse_pdf
takes a PDF, so we pack the individual page images into PDFs — each image
becomes one page, at native resolution (PyMuPDF embeds it, no re-encode).
A ``mapping.json`` records exactly which image landed on which page of which
PDF, so the per-PDF parser output can be split back to per-image results after
the run.
Output layout::
<out>/
batch_00000.pdf # 32 pages = 32 images
batch_00001.pdf
...
mapping.json
``mapping.json``::
{
"per_pdf": 32,
"num_images": 1651,
"num_pdfs": 52,
"pdfs": [
{"pdf": "batch_00000.pdf", "images": ["imgA.jpg", "imgB.jpg", ...]},
...
]
}
Page index (0-based) of an image == its position in that PDF's ``images`` list.
Example
-------
# chạy từ benchmark/parser/
python run_parser/build_pdfs.py \
--images data/images \
--out data/pdfs \
--per-pdf 32
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--images", required=True, type=Path,
help="Folder of page images (one image per page).")
parser.add_argument("--out", required=True, type=Path,
help="Output folder for the batched PDFs + mapping.json.")
parser.add_argument("--per-pdf", type=int, default=32,
help="Number of images packed into each PDF (default 32).")
parser.add_argument("--limit", type=int, default=None,
help="Only use the first N images (quick test).")
return parser.parse_args()
def find_images(images_dir: Path) -> list[Path]:
if not images_dir.is_dir():
raise NotADirectoryError(f"--images is not a folder: {images_dir}")
files = [p for p in sorted(images_dir.iterdir())
if p.is_file() and p.suffix.lower() in IMAGE_EXTENSIONS]
if not files:
raise FileNotFoundError(f"No image files found in {images_dir}")
return files
def image_to_pdf_pages(doc, image_path: Path) -> bool:
"""Append the image as one page to ``doc``. Returns True on success."""
import fitz # PyMuPDF
try:
with fitz.open(image_path) as img_doc:
pdf_bytes = img_doc.convert_to_pdf()
with fitz.open("pdf", pdf_bytes) as img_pdf:
doc.insert_pdf(img_pdf)
return True
except Exception as exc:
print(f" [skip] {image_path.name}: {exc!r}", flush=True)
return False
def main() -> int:
args = parse_args()
import fitz # PyMuPDF
images = find_images(args.images)
if args.limit is not None:
images = images[: args.limit]
args.out.mkdir(parents=True, exist_ok=True)
per_pdf = max(1, args.per_pdf)
num_pdfs = (len(images) + per_pdf - 1) // per_pdf
width = max(5, len(str(num_pdfs - 1)))
print(f"[build_pdfs] {len(images)} images -> {num_pdfs} PDFs "
f"({per_pdf} images/PDF) in {args.out}", flush=True)
pdf_entries: list[dict] = []
for pdf_idx in range(num_pdfs):
chunk = images[pdf_idx * per_pdf: (pdf_idx + 1) * per_pdf]
pdf_name = f"batch_{pdf_idx:0{width}d}.pdf"
out_pdf = args.out / pdf_name
doc = fitz.open()
used_images: list[str] = []
try:
for image_path in chunk:
if image_to_pdf_pages(doc, image_path):
used_images.append(image_path.name)
if len(doc) == 0:
print(f" [warn] {pdf_name}: no valid pages, skipped", flush=True)
continue
doc.save(out_pdf)
finally:
doc.close()
pdf_entries.append({"pdf": pdf_name, "images": used_images})
print(f" [{pdf_idx + 1}/{num_pdfs}] {pdf_name} {len(used_images)} pages",
flush=True)
mapping = {
"per_pdf": per_pdf,
"num_images": sum(len(e["images"]) for e in pdf_entries),
"num_pdfs": len(pdf_entries),
"pdfs": pdf_entries,
}
mapping_path = args.out / "mapping.json"
mapping_path.write_text(json.dumps(mapping, indent=2, ensure_ascii=False),
encoding="utf-8")
print(f"\n[build_pdfs] done. {mapping['num_pdfs']} PDFs, "
f"{mapping['num_images']} pages total.", flush=True)
print(f" mapping: {mapping_path}", flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|