inference.py

#!/usr/bin/env python3
"""Standalone runtime for the exported reader.

    python inference.py --model-dir . --image receipt.jpg \
        --question "가게 전화번호의 뒤에서 2번째 숫자는 무엇입니까?"

The graph runs once and produces a record. Questions are answered from that
record by `question_router`, so several questions about the same image cost one
forward pass -- pass `--question` more than once to see it.

Requires onnxruntime, numpy, and Pillow. No tokenizer, no vocabulary.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from question_router import (
    address_op_from_question, answer_from_record, phone_op_from_question,
    route_family_from_question,
)


def load_config(model_dir: Path) -> dict:
    return json.loads((model_dir / "config.json").read_text(encoding="utf-8"))


def preprocess(image_path: Path, cfg: dict) -> np.ndarray:
    from PIL import Image

    spec = cfg["input"]
    img = Image.open(image_path).convert("L").resize(
        (int(spec["width"]), int(spec["height"])), Image.Resampling.BILINEAR
    )
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return ((arr - 0.5) / 0.5)[None, None, :, :]


def decode(logits: np.ndarray, cfg: dict) -> dict:
    blank = int(cfg["blank_class"])
    phone_slots = int(cfg["phone_slots"])
    ids = logits[0].argmax(-1)

    def run(values) -> str:
        out = []
        for v in values:
            if int(v) == blank:
                break
            out.append(str(int(v)))
        return "".join(out)

    return {"phone": run(ids[:phone_slots]), "street": run(ids[phone_slots:])}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model-dir", default=".")
    ap.add_argument("--image", required=True)
    ap.add_argument("--question", action="append", default=[])
    ap.add_argument("--precision", choices=["fp32", "int8"], default="fp32")
    ap.add_argument("--threads", type=int, default=1)
    args = ap.parse_args()

    import onnxruntime as ort

    model_dir = Path(args.model_dir)
    cfg = load_config(model_dir)
    name = "model.onnx" if args.precision == "fp32" else "model_int8.onnx"
    path = model_dir / name
    if not path.exists():
        raise SystemExit(f"missing {path}")

    opts = ort.SessionOptions()
    opts.intra_op_num_threads = args.threads
    opts.inter_op_num_threads = args.threads
    sess = ort.InferenceSession(str(path), opts, providers=["CPUExecutionProvider"])

    logits = sess.run(None, {"image": preprocess(Path(args.image), cfg)})[0]
    record = decode(logits, cfg)

    answers = []
    for q in args.question:
        family = route_family_from_question(q)
        if family == "phone":
            op = phone_op_from_question(q)
        elif family == "address":
            op = address_op_from_question(q)
        else:
            op = "unsupported"
        answers.append({
            "question": q,
            "family": family,
            "op": op,
            "answer": answer_from_record(q, record["phone"], record["street"]),
        })

    print(json.dumps({
        "image": args.image,
        "precision": args.precision,
        "record": record,
        "forward_passes": 1,
        "answers": answers,
    }, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())