208 lines
10 KiB
Python
208 lines
10 KiB
Python
# Converts a PP-OCRv5 detector and recognizer to Core ML models for modules/OCR on macOS, then checks
|
|
# each against the ONNX model on the CPU. Run once, on a Mac, after modules/OCR/fetch_models.jai;
|
|
# mcap only uses the .mlpackage files it writes (into modules/OCR/models, beside the ONNX ones).
|
|
#
|
|
# 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
|
|
# ~/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
|
|
# the detector's text pixels and changes recognized text, and the Neural Engine gets it wrong.
|
|
#
|
|
# The shapes. Core ML on the GPU cannot take a flexible height and width for the detector (its
|
|
# depthwise convolutions fail to compile: "invalid axis" in mps.expand_dims), so both models take a
|
|
# list of sizes (at most 128 a model), which mcap pads its inputs up to; models.config's coreml_ keys
|
|
# 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 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
|
|
import onnx, onnxruntime as ort
|
|
from onnx import helper, numpy_helper
|
|
import torch, onnx2torch
|
|
import coremltools as ct
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
MODELS = os.path.join(HERE, "..", "modules", "OCR", "models")
|
|
|
|
DET_HEIGHTS = list(range(128, 1536 + 1, 128))
|
|
DET_WIDTHS = list(range(256, 2560 + 1, 256))
|
|
REC_BATCH = 16
|
|
REC_HEIGHT = 48
|
|
REC_WIDTHS = list(range(64, 4096 + 1, 64))
|
|
|
|
# coremltools 9.0 casts a folded one-element tensor with int(array), which NumPy 2 refuses ("only
|
|
# 0-dimensional arrays can be converted"); the same cast, with the element taken out first.
|
|
from coremltools.converters.mil.frontend.torch import ops as ct_ops
|
|
_original_cast = ct_ops._cast
|
|
def _cast_one_element(context, node, dtype, dtype_name):
|
|
x = ct_ops._get_inputs(context, node, expected=1)[0]
|
|
if x.can_be_folded_to_const() and np.ndim(x.val) > 0 and np.size(x.val) == 1:
|
|
context.add(ct_ops.mb.const(val=dtype(np.asarray(x.val).item()), name=node.name), node.name)
|
|
return
|
|
_original_cast(context, node, dtype, dtype_name)
|
|
ct_ops._cast = _cast_one_element
|
|
|
|
|
|
def constant_reshapes(model, input_shapes):
|
|
"""The model with each computed Reshape target made a constant of 0s and -1s."""
|
|
graph = model.graph
|
|
initializers = {i.name for i in graph.initializer}
|
|
producers = {o: n for n in graph.node for o in n.output}
|
|
computed = [n for n in graph.node if n.op_type == "Reshape" and n.input[1] not in initializers
|
|
and not (n.input[1] in producers and producers[n.input[1]].op_type == "Constant")]
|
|
if not computed:
|
|
return model
|
|
probe = onnx.ModelProto(); probe.CopyFrom(model)
|
|
names = []
|
|
for n in computed:
|
|
for t in (n.input[0], n.output[0]):
|
|
if t not in names:
|
|
names.append(t)
|
|
probe.graph.output.append(helper.make_tensor_value_info(t, onnx.TensorProto.FLOAT, None))
|
|
session = ort.InferenceSession(probe.SerializeToString(), providers=["CPUExecutionProvider"])
|
|
name = session.get_inputs()[0].name
|
|
shapes = []
|
|
for shape in input_shapes:
|
|
outs = session.run(names, {name: np.zeros(shape, np.float32)})
|
|
shapes.append({t: o.shape for t, o in zip(names, outs)})
|
|
for n in computed:
|
|
ins = [s[n.input[0]] for s in shapes]
|
|
outs = [s[n.output[0]] for s in shapes]
|
|
target = []
|
|
for d in range(len(outs[0])):
|
|
sizes = {o[d] for o in outs}
|
|
if len(sizes) == 1:
|
|
target.append(sizes.pop())
|
|
elif all(len(i) > d and i[d] == o[d] for i, o in zip(ins, outs)):
|
|
target.append(0)
|
|
else:
|
|
target.append(-1)
|
|
if target.count(-1) > 1:
|
|
sys.exit(f"{n.name}: more than one dimension to infer ({target})")
|
|
constant = n.name + "_constant_shape"
|
|
graph.initializer.append(numpy_helper.from_array(np.array(target, dtype=np.int64), constant))
|
|
n.input[1] = constant
|
|
print(f" {n.name}: target shape {target}")
|
|
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():
|
|
traced = torch.jit.trace(module, torch.from_numpy(example))
|
|
ml = ct.convert(traced, inputs=[ct.TensorType(name="x", shape=shape)], outputs=[ct.TensorType(name="y")],
|
|
convert_to="mlprogram", compute_precision=ct.precision.FLOAT32,
|
|
minimum_deployment_target=ct.target.macOS14)
|
|
ml.save(path)
|
|
return ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_GPU)
|
|
|
|
|
|
def ctc(logits):
|
|
texts = []
|
|
for row in logits.argmax(-1):
|
|
seq, last = [], -1
|
|
for i in row:
|
|
if i != last and i != 0:
|
|
seq.append(int(i))
|
|
last = i
|
|
texts.append(seq)
|
|
return texts
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", choices=["mobile", "server"], default="mobile")
|
|
args = parser.parse_args()
|
|
rng = np.random.default_rng(0)
|
|
|
|
# Detector.
|
|
stem = f"ch_PP-OCRv5_det_{args.model}"
|
|
onnx_path = os.path.join(MODELS, stem + ".onnx")
|
|
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()
|
|
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"])
|
|
worst = 0.0
|
|
for (h, w) in [(768, 1280), (1536, 2560), (128, 256)]:
|
|
x = rng.standard_normal((1, 3, h, w)).astype(np.float32)
|
|
a = reference.run(None, {"x": x})[0]
|
|
b = det.predict({"x": x})["y"]
|
|
flips = float(np.mean((a > 0.3) != (b > 0.3)))
|
|
worst = max(worst, flips)
|
|
print(f" {h} x {w}: max difference {np.abs(a - b).max():.2e}, pixels across 0.3: {flips * 100:.4f}%")
|
|
if worst > 0.0001:
|
|
sys.exit("The Core ML detector is too far from the ONNX one.")
|
|
|
|
# Recognizer.
|
|
stem = f"ch_PP-OCRv5_rec_{args.model}"
|
|
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(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),
|
|
ct.EnumeratedShapes(shapes=shapes, default=(REC_BATCH, 3, REC_HEIGHT, 640)), out_path)
|
|
print(f"Wrote {out_path} ({len(shapes)} widths) in {time.perf_counter() - start:.1f} s.")
|
|
reference = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
|
|
for w in (64, 640, 4096):
|
|
x = rng.standard_normal((REC_BATCH, 3, REC_HEIGHT, w)).astype(np.float32) * 0.5
|
|
a = reference.run(None, {"x": x})[0]
|
|
b = rec.predict({"x": x})["y"]
|
|
same = np.mean([p == q for p, q in zip(ctc(a), ctc(b))])
|
|
print(f" width {w}: max difference {np.abs(a - b).max():.2e}, decoded the same {same * 100:.0f}%")
|
|
if same < 1:
|
|
sys.exit("The Core ML recognizer reads differently from the ONNX one.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|