File size: 2,630 Bytes
6d05289 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | #!/usr/bin/env python3
"""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())
|