Add examples/batch.py

examples/batch.py

@@ -0,0 +1,81 @@
+"""Batch extraction over a directory of table-cropped images.
+
+Looks for ``<stem>.md`` alongside each image; if found, passes its contents
+as the ``markdown`` disambiguator. Without markdown the model picks an
+arbitrary plausible table and quality drops significantly.
+
+Usage:
+    python examples/batch.py <image_dir> [--preset quality|fast] [--out results.jsonl]
+"""
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from inference import PRESETS, StructuredExtractor
+
+CHECKPOINT = "Glazkov/structured-extractor-qwen3vl-4b-exp93"
+IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tif", ".tiff"}
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("image_dir", type=Path, help="Directory with table-crop images")
+    p.add_argument(
+        "--preset",
+        choices=list(PRESETS),
+        default="quality",
+        help="quality (beam=4 + rp=1.1), balanced (beam=4), fast (greedy). Default: quality.",
+    )
+    p.add_argument(
+        "--out",
+        type=Path,
+        default=Path("results.jsonl"),
+        help="Path to JSONL output file (one record per image).",
+    )
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    paths = sorted(p for p in args.image_dir.iterdir() if p.suffix.lower() in IMAGE_EXTS)
+    if not paths:
+        sys.exit(f"No images found in {args.image_dir}")
+
+    markdowns: list[str | None] = []
+    for p in paths:
+        md_path = p.with_suffix(".md")
+        markdowns.append(md_path.read_text(encoding="utf-8") if md_path.exists() else None)
+    print(f"Loaded {sum(1 for m in markdowns if m)}/{len(paths)} markdown sidecars.")
+
+    print(f"Loading checkpoint {CHECKPOINT} ...")
+    extractor = StructuredExtractor.from_pretrained(CHECKPOINT)
+
+    print(f"Running '{args.preset}' over {len(paths)} images ...")
+    t0 = time.time()
+    results = extractor.extract_batch(
+        paths,
+        preset=args.preset,
+        markdown_batch=markdowns,
+        batch_size=1,
+        return_raw=False,
+    )
+    elapsed = time.time() - t0
+
+    with args.out.open("w", encoding="utf-8") as f:
+        for path, result in zip(paths, results):
+            record = {"image": str(path), "parameters": result["parameters"]}
+            f.write(json.dumps(record, ensure_ascii=False) + "\n")
+
+    print(
+        f"\nDone in {elapsed:.1f}s ({elapsed / len(paths):.2f}s/image). "
+        f"Wrote {len(paths)} records to {args.out}"
+    )
+
+
+if __name__ == "__main__":
+    main()