example.py

#!/usr/bin/env python3
"""Minimal honmono-ocr recognition inference — onnxruntime + CTC decode.

Transcribes a single Japanese text-line crop. For a full page (detection + this
recognizer), see demo/app.py.

Usage:
    python example.py [line_crop.jpg]

Paths auto-resolve for both layouts:
  * Hugging Face model repo (flat):  book_rec_fp16.onnx, jp_dict.txt
  * this GitHub repo:                models/book_rec_fp16.onnx,
                                     book_ocr/assets/jp_dict.txt, examples/line_sample.jpg

Note: rotate tall (vertical) line crops 90° CCW before recognition — see README.
On desktop ONNX Runtime the FP16 model needs ORT_ENABLE_BASIC (set below).
"""

import sys
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort


def _first_existing(*candidates: str) -> str:
    for c in candidates:
        if Path(c).exists():
            return c
    return candidates[0]


MODEL = _first_existing("book_rec_fp16.onnx", "models/book_rec_fp16.onnx")
DICT = _first_existing("jp_dict.txt", "book_ocr/assets/jp_dict.txt")


def load_chars(path: str) -> list[str]:
    """['blank'] + dict chars (file order) + [' ']  (use_space_char=True).

    Strip only the trailing newline — the dict includes real whitespace glyphs
    (U+0020 at line 0, U+3000 at line 327); dropping them shifts every index.
    """
    chars = ["blank"]
    with open(path, encoding="utf-8") as f:
        for ln in f:
            ch = ln.rstrip("\r\n")
            if ch:
                chars.append(ch)
    return chars + [" "]


def preprocess(bgr: np.ndarray) -> np.ndarray:
    """BGR crop -> [1,3,48,480] float32, normalized to [-1,1], right-padded."""
    h, w = bgr.shape[:2]
    new_w = min(max(int(round(w * 48 / h)), 1), 480)
    img = cv2.resize(bgr, (new_w, 48)).astype(np.float32).transpose(2, 0, 1) / 255.0
    img = (img - 0.5) / 0.5
    out = np.zeros((3, 48, 480), np.float32)
    out[:, :, :new_w] = img
    return out[None]


def ctc_decode(logits: np.ndarray, chars: list[str]) -> str:
    """Greedy CTC: argmax per timestep, drop blanks (0), collapse repeats."""
    prev, res = -1, []
    for i in logits.argmax(-1):
        if i != 0 and i != prev and i < len(chars):
            res.append(chars[i])
        prev = i
    return "".join(res)


def main() -> None:
    image = sys.argv[1] if len(sys.argv) > 1 else _first_existing("line.jpg", "examples/line_sample.jpg")
    crop = cv2.imread(image)
    if crop is None:
        raise SystemExit(f"could not read image: {image}")

    opts = ort.SessionOptions()
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC  # FP16 desktop-ORT fix
    sess = ort.InferenceSession(MODEL, opts, providers=["CPUExecutionProvider"])
    chars = load_chars(DICT)

    logits = sess.run(None, {sess.get_inputs()[0].name: preprocess(crop)})[0][0]
    print(ctc_decode(logits, chars))


if __name__ == "__main__":
    main()