| |
| """Logical-order Persian inference for Bina 0.2 - RizehPizeh.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import re |
| from pathlib import Path |
| from typing import Any, Iterator |
|
|
| from paddleocr import TextRecognition |
|
|
|
|
| _LTR_RUN = re.compile(r"[a-zA-Z0-9 :*./%+-]") |
|
|
|
|
| def pred_reverse(text: str) -> str: |
| """Match PaddleOCR's Arabic-aware BaseRecLabelDecode.pred_reverse.""" |
| segments: list[str] = [] |
| current_ltr = "" |
| for character in text: |
| if _LTR_RUN.search(character): |
| current_ltr += character |
| continue |
| if current_ltr: |
| segments.append(current_ltr) |
| current_ltr = "" |
| segments.append(character) |
| if current_ltr: |
| segments.append(current_ltr) |
| return "".join(reversed(segments)) |
|
|
|
|
| class BinaTextRecognition: |
| """Run the exported Paddle model and return logical-order Persian text.""" |
|
|
| def __init__( |
| self, |
| model_dir: str | Path | None = None, |
| device: str | None = None, |
| ) -> None: |
| if model_dir is None: |
| model_dir = Path(__file__).resolve().parent / "inference" |
| options: dict[str, Any] = {"model_dir": str(model_dir)} |
| if device is not None: |
| options["device"] = device |
| self._model = TextRecognition(**options) |
|
|
| def predict( |
| self, |
| input: str | Path | list[str] | list[Path], |
| batch_size: int = 1, |
| ) -> Iterator[dict[str, Any]]: |
| for result in self._model.predict(input=input, batch_size=batch_size): |
| payload = result.json |
| if callable(payload): |
| payload = payload() |
| raw = payload["res"] |
| visual_text = str(raw["rec_text"]) |
| yield { |
| "input_path": raw.get("input_path"), |
| "text": pred_reverse(visual_text), |
| "score": float(raw["rec_score"]), |
| "raw_visual_text": visual_text, |
| } |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("images", nargs="+") |
| parser.add_argument( |
| "--model-dir", |
| default=str(Path(__file__).resolve().parent / "inference"), |
| ) |
| parser.add_argument("--device", default=None) |
| parser.add_argument("--batch-size", type=int, default=1) |
| args = parser.parse_args() |
|
|
| model = BinaTextRecognition(args.model_dir, device=args.device) |
| for prediction in model.predict(args.images, batch_size=args.batch_size): |
| print(json.dumps(prediction, ensure_ascii=False)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
