Spaces:
Running on Zero
Running on Zero
| """Export original page pixels losslessly, with per-document provenance.""" | |
| import json | |
| from pathlib import Path | |
| from time import perf_counter | |
| from visual_doc_rag.data import load_tables, write_json | |
| def export_pages(subset: str, data_dir: Path) -> dict: | |
| started = perf_counter() | |
| source, tables = load_tables(subset, data_dir) | |
| root = data_dir / "pages" / subset / source["revision"] | |
| root.mkdir(parents=True, exist_ok=True) | |
| metadata = {row["doc_id"]: row for row in tables["documents_metadata"]} | |
| pages = [] | |
| for index, row in enumerate(tables["corpus"]): | |
| path = root / f"{row['corpus_id']}.png" | |
| with row["image"] as image: | |
| if not path.exists(): | |
| temporary = path.with_suffix(".tmp") | |
| image.convert("RGB").save(temporary, format="PNG") | |
| temporary.replace(path) | |
| size = list(image.size) | |
| document = metadata[row["doc_id"]] | |
| pages.append( | |
| { | |
| "corpus_id": row["corpus_id"], | |
| "doc_id": row["doc_id"], | |
| "page_number_in_doc": row["page_number_in_doc"], | |
| "file": path.name, | |
| "size": size, | |
| "license": document["license"], | |
| "source_url": document["url"], | |
| "source_file_name": document["file_name"], | |
| } | |
| ) | |
| if (index + 1) % 100 == 0: | |
| print(f"Pages: {index + 1}/{len(tables['corpus'])}", flush=True) | |
| manifest = { | |
| "dataset": source, | |
| "pages": pages, | |
| "format": "PNG, original pixels, RGB", | |
| "export_seconds": perf_counter() - started, | |
| "image_bytes": sum((root / row["file"]).stat().st_size for row in pages), | |
| } | |
| write_json(root / "manifest.json", manifest) | |
| return manifest | |
| def load_pages(subset: str, data_dir: Path, revision: str) -> tuple[Path, dict]: | |
| root = data_dir / "pages" / subset / revision | |
| path = root / "manifest.json" | |
| if not path.exists(): | |
| raise ValueError("Page cache missing for this dataset revision. Run `vdr pages` first.") | |
| return root, json.loads(path.read_text(encoding="utf-8")) | |