[0.1.31] macOS phase 5: PP-OCRv5 server OCR on the GPU through native Core ML

README.md

@@ -30,8 +30,9 @@ see [Screen capture](#screen-capture).
 
    **On macOS**, labelling and OCR run on the GPU through Core ML, from Core ML copies of the models
    made once with `tools/export_siglip_coreml.py` (SigLIP 2's image model, about ten times as fast as
-   the CPU) and `tools/export_ppocr_coreml.py` (PP-OCRv5 mobile's detector and recognizer, about twice
-   as fast; their first lines say how to set up the Python they need). Without them, each runs on the
+   the CPU) and `tools/export_ppocr_coreml.py` (PP-OCRv5's detector and recognizer, about twice as
+   fast; run it with `--model mobile` and `--model server`; their first lines say how to set up the
+   Python they need). Without them, each runs on the
    CPU. The first scan after making them is slower: Core ML prepares each image size the first time it
    meets it, and macOS keeps that for later runs.
 

check_targets.jai

@@ -49,7 +49,7 @@ APP_NAME        :: "mcap";
 APP_FOLDER_NAME :: "mcap";
 GIT_BRANCH    :: "check";
 GIT_HASH      :: "-";
-VERSION       :: "0.1.30";
+VERSION       :: "0.1.31";
 BUILD_DATE    :: "-";
 DEBUG         :: true;
 COMPILE_TESTS :: true;

first.jai

@@ -8,7 +8,7 @@ APP_NAME              :: "mcap";
 APP_FOLDER_NAME       :: "mcap";
 APP_BUNDLE_IDENTIFIER :: "com.example.mcap";
 
-VERSION :: "0.1.30";
+VERSION :: "0.1.31";
 // Where everything goes: bin/, or bin-headless/ for builds made with no_crash_dialog, so automated
 // builds never write over the ones you run by hand.
 build_directory := "bin";

modules/OCR/models.config

@@ -54,3 +54,5 @@ unclip_ratio:    1.5
 rec_height:      48
 rec_batch_gpu:   8
 rec_batch_cpu:   1
+detector_coreml:   "ch_PP-OCRv5_det_server.mlpackage"
+recognizer_coreml: "ch_PP-OCRv5_rec_server.mlpackage"

tools/export_ppocr_coreml.py

@@ -5,7 +5,8 @@
 #   python3 -m venv ~/mcap-coreml
 #   ~/mcap-coreml/bin/pip install torch==2.7.0 coremltools==9.0 onnx onnxruntime pillow
 #   ~/mcap-coreml/bin/pip install --no-deps onnx2torch torchvision==0.22.0
-#   ~/mcap-coreml/bin/python tools/export_ppocr_coreml.py [--model mobile|server]
+#   ~/mcap-coreml/bin/python tools/export_ppocr_coreml.py --model mobile
+#   ~/mcap-coreml/bin/python tools/export_ppocr_coreml.py --model server
 #
 # There is no PyTorch source for PP-OCR (it is PaddlePaddle's), and coremltools no longer reads ONNX, so
 # each model goes ONNX -> PyTorch (onnx2torch) -> Core ML, as an ML Program in float32: float16 moves
@@ -17,12 +18,17 @@
 # say the same steps, and must be kept in step with these:
 #   detector:   1 x 3 x H x W, H a multiple of 128 up to 1536 (det_tile_height), W of 256 up to 2560 (det_max_side)
 #   recognizer: 16 x 3 x 48 x W, W a multiple of 64 up to 4096 (REC_MAX_WIDTH); fewer lines are padded to 16
-# The recognizer's Reshapes compute their target shape from tensor sizes, which tracing freezes at the
-# traced width; they are rewritten first as constant shapes of 0 (copy the dimension) and -1 (infer it),
-# worked out from the shapes onnxruntime gives at several sizes, and checked to give the same output.
+# The server models pad some 2 x 2 convolutions and max pools with auto_pad SAME_UPPER, which onnx2torch
+# does not take; at stride 1 it is a fixed pixel after, written out as pads first.
+# Reshapes that compute their target shape from tensor sizes (the recognizers', the server detector's)
+# are frozen by tracing at the traced size; they are rewritten first as constant shapes of 0 (copy the
+# dimension) and -1 (infer it), worked out from the shapes onnxruntime gives at several sizes, and
+# checked to give the same output.
 #
 # Measured on an M5 Pro against onnxruntime on the CPU, each at the same (padded) size: the detector
 # about 5 times as fast, the recognizer about 4 times for 16 lines, with the same text pixels and text.
+# Whole images through mcap (ocr-bench, 8 text pages): mobile 323 ms an image against 536, server 940
+# against 2030.
 
 import argparse, itertools, os, sys, time
 import numpy as np
@@ -96,6 +102,34 @@ def constant_reshapes(model, input_shapes):
     return model
 
 
+def explicit_padding(model):
+    """The model with auto_pad SAME_UPPER and SAME_LOWER, which onnx2torch does not take, written out as
+    pads. Only stride 1 is rewritten: there the padding does not depend on the input's size (the server
+    models' 2 x 2 convolutions and max pools: one pixel after, none before)."""
+    for n in model.graph.node:
+        attrs = {a.name: a for a in n.attribute}
+        mode = attrs["auto_pad"].s.decode() if "auto_pad" in attrs else "NOTSET"
+        if mode not in ("SAME_UPPER", "SAME_LOWER"):
+            continue
+        kernel = list(attrs["kernel_shape"].ints) if "kernel_shape" in attrs else None
+        if kernel is None:
+            sys.exit(f"{n.name}: auto_pad {mode} with no kernel_shape")
+        strides = list(attrs["strides"].ints) if "strides" in attrs else [1] * len(kernel)
+        dilations = list(attrs["dilations"].ints) if "dilations" in attrs else [1] * len(kernel)
+        if any(s != 1 for s in strides):
+            sys.exit(f"{n.name}: auto_pad {mode} with strides {strides}; only stride 1 is rewritten")
+        totals = [(k - 1) * d for k, d in zip(kernel, dilations)]
+        small = [t // 2 for t in totals]
+        begins = small if mode == "SAME_UPPER" else [t - b for t, b in zip(totals, small)]
+        ends = [t - b for t, b in zip(totals, begins)]
+        n.attribute.remove(attrs["auto_pad"])
+        if "pads" in attrs:
+            n.attribute.remove(attrs["pads"])
+        n.attribute.append(helper.make_attribute("pads", begins + ends))
+        print(f"  {n.name} ({n.op_type}): auto_pad {mode} -> pads {begins + ends}")
+    return model
+
+
 def convert(onnx_model, example, shape, path):
     module = onnx2torch.convert(onnx_model).eval()
     with torch.no_grad():
@@ -131,7 +165,8 @@ def main():
     out_path = os.path.join(MODELS, stem + ".mlpackage")
     shapes = [(1, 3, h, w) for h in DET_HEIGHTS for w in DET_WIDTHS]
     start = time.perf_counter()
-    det = convert(onnx.load(onnx_path), np.zeros((1, 3, 768, 1280), np.float32),
+    model = constant_reshapes(explicit_padding(onnx.load(onnx_path)), [(1, 3, h, w) for h in (256, 512) for w in (512, 768)])
+    det = convert(model, np.zeros((1, 3, 768, 1280), np.float32),
                   ct.EnumeratedShapes(shapes=shapes, default=(1, 3, 768, 1280)), out_path)
     print(f"Wrote {out_path} ({len(shapes)} sizes) in {time.perf_counter() - start:.1f} s.")
     reference = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
@@ -151,7 +186,7 @@ def main():
     onnx_path = os.path.join(MODELS, stem + ".onnx")
     out_path = os.path.join(MODELS, stem + ".mlpackage")
     print("Recognizer: making its reshapes constant")
-    model = constant_reshapes(onnx.load(onnx_path), [(n, 3, REC_HEIGHT, w) for n in (1, 2) for w in (320, 640)])
+    model = constant_reshapes(explicit_padding(onnx.load(onnx_path)), [(n, 3, REC_HEIGHT, w) for n in (1, 2) for w in (320, 640)])
     shapes = [(REC_BATCH, 3, REC_HEIGHT, w) for w in REC_WIDTHS]
     start = time.perf_counter()
     rec = convert(model, np.zeros((REC_BATCH, 3, REC_HEIGHT, 640), np.float32),