| |
| """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()) |
|
|
