File size: 1,952 Bytes
4dce182 | 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 | """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()
|