File size: 5,409 Bytes
139ea68 | 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 104 105 106 107 | """Reference of the browser OCR pipeline (src/lib/docai/ocr.js) on PP-OCRv5 mobile det + en_PP-OCRv5 mobile rec.
det: longest side resized to LIMIT (multiple of 32), BGR, ImageNet mean/std -> probability map;
threshold 0.3 -> connected components -> axis-aligned boxes scored by mean probability (>= 0.6),
grown by area * 1.5 / perimeter (DB's unclip for a rectangle)
rec: each box cropped, height 48, width bucketed, (x/255 - 0.5) / 0.5, BGR -> CTC greedy decode;
each kept character's time step gives its x position, so lines split into word boxes.
Run with .models/.venv311mp (opencv, onnxruntime). usage: python ocr_reference.py [image]
"""
import os, sys, time
import numpy as np
import cv2
import onnxruntime as ort
from PIL import Image, ImageDraw, ImageFont
HERE = os.path.dirname(os.path.abspath(__file__))
PORT = os.path.join(HERE, "..", "portfolio", "pp-ocrv5")
DICT = open(os.path.join(PORT, "en_dict.txt"), encoding="utf-8").read().split("\n")[:-1] + [" "]
LIMIT, THRESH, BOX_THRESH, UNCLIP = 1280, 0.3, 0.6, 1.5
MEAN, STD = np.array([0.485, 0.456, 0.406], np.float32), np.array([0.229, 0.224, 0.225], np.float32)
def synthetic_invoice():
img = Image.new("RGB", (1240, 1600), "white")
d = ImageDraw.Draw(img)
f = lambda s, b=False: ImageFont.truetype("C:/Windows/Fonts/arialbd.ttf" if b else "C:/Windows/Fonts/arial.ttf", s)
d.text((80, 80), "ACME Supplies Ltd.", font=f(44, True), fill="black")
d.text((900, 90), "INVOICE", font=f(52, True), fill=(240, 72, 44))
d.text((80, 150), "221B Baker Street, London NW1 6XE", font=f(24), fill="#333")
for i, (k, v) in enumerate([("Invoice #:", "INV-2024-0412"), ("Date:", "12 March 2024"), ("Due:", "11 April 2024")]):
d.text((820, 200 + i * 36), k, font=f(24, True), fill="black")
d.text((960, 200 + i * 36), v, font=f(24), fill="black")
d.text((80, 360), "Description", font=f(26, True), fill="black")
d.text((900, 360), "Amount", font=f(26, True), fill="black")
for i, (k, v) in enumerate([("Printer paper A4 (10 reams)", "$45.90"), ("Toner cartridge HP 26A", "$129.00"), ("Delivery", "$12.50")]):
d.text((80, 420 + i * 48), k, font=f(26), fill="black")
d.text((900, 420 + i * 48), v, font=f(26), fill="black")
d.text((700, 620), "Total due: $187.40", font=f(34, True), fill="black")
d.text((80, 760), "Thank you for your business! Questions? billing@acme.co.uk", font=f(20), fill="#555")
return np.array(img)[:, :, ::-1] # BGR
def detect(sess, bgr):
h, w = bgr.shape[:2]
s = min(1.0, LIMIT / max(h, w)) if max(h, w) > LIMIT else LIMIT / max(h, w)
rh, rw = max(32, int(round(h * s / 32)) * 32), max(32, int(round(w * s / 32)) * 32)
x = cv2.resize(bgr, (rw, rh)).astype(np.float32) / 255
x = ((x - MEAN) / STD).transpose(2, 0, 1)[None]
prob = sess.run(None, {"x": x})[0][0, 0]
n, labels, stats, _ = cv2.connectedComponentsWithStats((prob > THRESH).astype(np.uint8), connectivity=8)
boxes = []
for i in range(1, n):
x0, y0, bw, bh, area = stats[i]
if min(bw, bh) < 3:
continue
score = float(prob[labels == i].mean())
if score < BOX_THRESH:
continue
d = bw * bh * UNCLIP / (2 * (bw + bh))
sx, sy = w / rw, h / rh
boxes.append([max(0, (x0 - d) * sx), max(0, (y0 - d) * sy), min(w, (x0 + bw + d) * sx), min(h, (y0 + bh + d) * sy), score])
boxes.sort(key=lambda b: (round(b[1] / 20), b[0]))
return boxes
def recognize(sess, bgr, boxes):
lines = []
for x0, y0, x1, y1, _ in boxes:
crop = bgr[int(y0):int(np.ceil(y1)), int(x0):int(np.ceil(x1))]
ch, cw = crop.shape[:2]
tw = max(16, int(np.ceil(48 * cw / ch)))
x = cv2.resize(crop, (tw, 48)).astype(np.float32) / 255
x = ((x - 0.5) / 0.5).transpose(2, 0, 1)[None]
probs = sess.run(None, {"x": x})[0][0] # [T, 438]
T = probs.shape[0]
ids, conf = probs.argmax(1), probs.max(1)
chars, prev = [], 0
for t, k in enumerate(ids):
if k != 0 and k != prev:
chars.append((DICT[k - 1], t, conf[t]))
prev = k
text = "".join(c for c, _, _ in chars)
# words from the character time steps
words, cur = [], []
for c in chars + [(" ", T, 1.0)]:
if c[0] == " ":
if cur:
t0, t1 = cur[0][1], cur[-1][1]
words.append(("".join(k for k, _, _ in cur), x0 + (t0 - 0.5) * cw / T, x0 + (t1 + 1.5) * cw / T, float(np.mean([p for _, _, p in cur]))))
cur = []
else:
cur.append(c)
lines.append((text, float(np.mean([p for _, _, p in chars])) if chars else 0, (x0, y0, x1, y1), words))
return lines
if __name__ == "__main__":
bgr = cv2.imread(sys.argv[1]) if len(sys.argv) > 1 else synthetic_invoice()
det = ort.InferenceSession(os.path.join(PORT, "det.onnx"), providers=["CPUExecutionProvider"])
rec = ort.InferenceSession(os.path.join(PORT, "rec.onnx"), providers=["CPUExecutionProvider"])
t0 = time.time()
boxes = detect(det, bgr)
lines = recognize(rec, bgr, boxes)
print(f"{len(lines)} lines in {time.time() - t0:.2f}s")
for text, conf, box, words in lines:
print(f"{conf:.2f} {[int(v) for v in box]} {text!r} | words: {[(w[0], int(w[1]), int(w[2])) for w in words]}")
|