File size: 3,842 Bytes
1e4f05f febacad | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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())
|