Add sd-turbo-webgpu-q4, nanodiffuse-q8, whisper-tiny-timestamped, whisper-base-timestamped, distilbert-sst2, distilbert-go-emotions, distilbert-ner, rtdetr-v2-r18, depth-anything-v2-small, slimsam-77, magenta-style, pp-ocrv5
139ea68 verified | """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]}") | |