| """tinydoc command line. |
| |
| python -m tinydoc extract receipt.jpg invoice.pdf --backend onnx --weights models |
| python -m tinydoc params |
| """ |
|
|
| import argparse |
| import json |
| import sys |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(prog="tinydoc") |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| ex = sub.add_parser("extract", help="extract structured JSON from images / PDFs") |
| ex.add_argument("files", nargs="+") |
| ex.add_argument("--backend", choices=["onnx", "torch"], default="onnx") |
| ex.add_argument("--weights", default=None, help="models/ (onnx) or runs/ (torch)") |
| ex.add_argument("--device", default="cpu") |
| ex.add_argument("--force-ocr", action="store_true", help="OCR PDFs even when they contain usable text") |
| ex.add_argument("--out", default=None, help="write JSON here (default: stdout)") |
| ex.add_argument("--compact", action="store_true", help="omit per-segment detail") |
| sub.add_parser("params", help="print the parameter budget") |
| args = ap.parse_args() |
|
|
| if args.cmd == "params": |
| from .models import BUDGET, parameter_report |
| for k, v in parameter_report().items(): |
| print(f"{k:12s} {v:>10,d} (target {BUDGET.get(k, sum(BUDGET.values())):,d})") |
| return |
|
|
| from .pipeline.engine import Engine |
| weights = args.weights or ("models" if args.backend == "onnx" else "runs") |
| engine = Engine(weights, backend=args.backend, device=args.device) |
| results = [] |
| for f in args.files: |
| r = engine.process(f, force_ocr=args.force_ocr) |
| if args.compact: |
| for p in r["pages"]: |
| p.pop("segments", None) |
| results.append(r) |
| payload = results[0] if len(results) == 1 else results |
| text = json.dumps(payload, indent=2, ensure_ascii=False) |
| if args.out: |
| with open(args.out, "w") as fh: |
| fh.write(text) |
| else: |
| sys.stdout.write(text + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
