cli.py

#!/usr/bin/env python3
"""
cli.py — batch-process a folder of documents from the command line.

Usage:
    python cli.py --input ./samples --output ./output
    python cli.py --input ./samples --output ./output --lang german
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

from app.detector import StampDetector

SUPPORTED = {".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".tif", ".bmp"}


def main() -> None:
    parser = argparse.ArgumentParser(description="Batch stamp detection")
    parser.add_argument("--input",  "-i", required=True, help="Input folder with PDFs / images")
    parser.add_argument("--output", "-o", default="./output", help="Output folder for JSON results")
    parser.add_argument("--lang",   "-l", default="german",  help="OCR language (default: german)")
    args = parser.parse_args()

    input_dir  = Path(args.input)
    output_dir = Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    files = [f for f in sorted(input_dir.iterdir()) if f.suffix.lower() in SUPPORTED]
    if not files:
        print(f"No supported files found in {input_dir}", file=sys.stderr)
        sys.exit(1)

    print(f"Found {len(files)} file(s). Loading model …")
    detector = StampDetector(lang=args.lang)
    print("Model ready. Processing …\n")

    summary = []
    for path in files:
        t0 = time.perf_counter()
        data = path.read_bytes()
        ext  = path.suffix.lower()

        if ext == ".pdf":
            result = detector.process_pdf_bytes(data, path.name)
        else:
            result = detector.process_image_bytes(data, path.name)

        elapsed = time.perf_counter() - t0
        stamp_count = result.get("total_stamps", result.get("stamp_count", 0))

        out_file = output_dir / (path.stem + ".json")
        out_file.write_text(json.dumps(result, ensure_ascii=False, indent=2))

        line = f"  {path.name:<40}  stamps: {stamp_count:>3}  ({elapsed:.2f}s)  → {out_file.name}"
        print(line)
        summary.append({"file": path.name, "stamps": stamp_count, "elapsed_s": round(elapsed, 2)})

    summary_file = output_dir / "_summary.json"
    summary_file.write_text(json.dumps(summary, ensure_ascii=False, indent=2))
    total = sum(s["stamps"] for s in summary)
    print(f"\nDone. {len(files)} file(s), {total} stamp(s) total. Summary → {summary_file}")


if __name__ == "__main__":
    main()