tools/test_quick.py

#!/usr/bin/env python3
"""
Quick OCR test — verify the ONNX engine works on sample images.

Usage:
    python tools/test_quick.py                          # test all images in working_space/input/
    python tools/test_quick.py path/to/image.png        # test single image
"""

import sys
import time
from pathlib import Path

from PIL import Image

sys.path.insert(0, str(Path(__file__).parent.parent))
from ocr.engine_onnx import OcrEngineOnnx


def test_image(engine: OcrEngineOnnx, path: Path) -> bool:
    """Run OCR on one image and print results."""
    img = Image.open(path)
    t0 = time.perf_counter()
    result = engine.recognize_pil(img)
    elapsed = (time.perf_counter() - t0) * 1000

    status = "OK" if result.text.strip() else "EMPTY"
    text_short = result.text.replace("\n", " | ")[:80]

    print(f"  [{status}] {path.name:30s} {elapsed:6.0f}ms  "
          f"L={len(result.lines):2d}  C={result.average_confidence:.0%}  "
          f'"{text_short}"')
    return bool(result.text.strip())


def main():
    if len(sys.argv) > 1:
        paths = [Path(p) for p in sys.argv[1:]]
    else:
        input_dir = Path("working_space/input")
        if not input_dir.exists():
            print(f"No images found. Place .png files in {input_dir}")
            return
        paths = sorted(input_dir.glob("*.png"))

    if not paths:
        print("No images to test.")
        return

    print(f"Testing {len(paths)} image(s)...\n")
    engine = OcrEngineOnnx()

    ok = sum(test_image(engine, p) for p in paths)
    print(f"\n  Result: {ok}/{len(paths)} images produced text")


if __name__ == "__main__":
    main()