example_usage.py

"""Minimal end-to-end use of the PP-OCRv6 Hanzi Unicode OCR model.

    python example_usage.py path/to/glyph.png

With no argument the script renders a few characters from a local font and
recognises them, so it can be pasted into a fresh environment as a smoke test.
"""
import json
import math
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont

MODEL_DIR = Path(__file__).resolve().parent

# Default fonts used only by the self-test path. Any TTF/OTF works.
FONT_MAIN = "C:/Windows/Fonts/simsun.ttc"
FONT_EXT_B = "C:/Windows/Fonts/simsunb.ttf"


class UnicodeOCR:
    def __init__(self, model_dir=MODEL_DIR, providers=None):
        model_dir = Path(model_dir)
        meta = json.loads((model_dir / "metadata.json").read_text(encoding="utf-8"))
        self.vocab = meta["vocab"]
        self.channels, self.height, self.width = meta["image_shape"]
        self.session = ort.InferenceSession(
            str(model_dir / "model.onnx"),
            providers=providers or ["CPUExecutionProvider"],
        )
        self.input_name = self.session.get_inputs()[0].name

    def prepare(self, image):
        """RGB -> BGR, aspect-preserving resize to height 48, [-1, 1], right pad."""
        image = image.convert("RGB")
        inner_w = min(math.ceil(self.height * image.width / image.height), self.width)
        bgr = np.asarray(image)[:, :, ::-1]
        resized = Image.fromarray(bgr).resize((inner_w, self.height), Image.BILINEAR)
        plane = np.asarray(resized).astype(np.float32).transpose(2, 0, 1)
        plane = plane / 127.5 - 1.0
        canvas = np.zeros((self.channels, self.height, self.width), np.float32)
        canvas[:, :, :inner_w] = plane
        return canvas[None]

    def recognise(self, image):
        """Return (text, mean_score) for one single-glyph image."""
        probs = self.session.run(None, {self.input_name: self.prepare(image)})[0][0]
        index = probs.argmax(-1)
        scores = probs.max(-1)
        text, kept, last = [], [], None
        for i, score in zip(index, scores):
            if i != last and i != 0:
                text.append(self.vocab[i])
                kept.append(float(score))
            last = i
        return "".join(text), (sum(kept) / len(kept) if kept else 0.0)


def render_glyph(char, font_path, px=200):
    font = ImageFont.truetype(font_path, px)
    left, top, right, bottom = font.getbbox(char)
    image = Image.new("L", (right - left + 20, bottom - top + 20), 255)
    ImageDraw.Draw(image).text((10 - left, 10 - top), char, font=font, fill=0)
    return image


def pick_font(char):
    # Characters above the BMP need a font that actually contains them, otherwise
    # the renderer substitutes .notdef and the model is asked to read a box.
    return FONT_EXT_B if ord(char) >= 0x20000 else FONT_MAIN


def self_test(model):
    chars = ["\u4e00", "\u54c1", "\u5b57", "\u660e", "\u6728", "\u6797", "\u68ee",
             "\u4dbf", "\u9fff", "\U00020000", "\U0002a6d6", "\U0002a6df"]
    hits = 0
    print(f"{'char':<6}{'codepoint':<12}{'predicted':<10}{'score':<10}ok")
    for char in chars:
        predicted, score = model.recognise(render_glyph(char, pick_font(char)))
        hit = predicted == char
        hits += hit
        print(f"{char:<6}{'U+%X' % ord(char):<12}{predicted:<10}{score:<10.4f}{hit}")
    print(f"\nmatched {hits}/{len(chars)}")
    return hits == len(chars)


def main():
    model = UnicodeOCR()
    if len(sys.argv) > 1:
        path = Path(sys.argv[1])
        text, score = model.recognise(Image.open(path))
        print(json.dumps({"file": str(path), "text": text, "score": score},
                         ensure_ascii=False))
        return 0
    return 0 if self_test(model) else 1


if __name__ == "__main__":
    sys.exit(main())