File size: 1,690 Bytes
ac871ba | 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 | #!/usr/bin/env python3
"""Command-line and importable entry point for Flippd pinball score extraction."""
from __future__ import annotations
import argparse
import json
from pathlib import Path
from typing import Sequence
from pinball_score_ocr import ScoreSuggester
class PinballScoreExtractor:
"""Load the packaged detector and CTC reader and extract score suggestions."""
def __init__(self, model_dir: str | Path | None = None, device: str | None = None):
root = Path(__file__).resolve().parent
self._suggester = ScoreSuggester(model_dir or root, device=device)
def predict(self, image: str | Path, confidence: float | None = None) -> dict:
"""Return ranked suggestions, detections, and stage timings for an image."""
return self._suggester.infer(image, bar=confidence)
def _parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Extract pinball scores from a display photo")
parser.add_argument("image", type=Path, help="Pillow-readable input image")
parser.add_argument("--model-dir", type=Path, help="directory containing config.json and both ONNX models")
parser.add_argument("--device", help="cpu, cuda, or cuda:N")
parser.add_argument("--confidence", type=float, help="minimum decoded suggestion confidence (default: config.json)")
return parser
def main(argv: Sequence[str] | None = None) -> int:
args = _parser().parse_args(argv)
extractor = PinballScoreExtractor(args.model_dir, args.device)
result = extractor.predict(args.image, args.confidence)
print(json.dumps(result, indent=2))
return 0
if __name__ == "__main__":
raise SystemExit(main())
|