| """ONNX (RT-DETRv4-S / D-FINE, ogkalu/comic-text-and-bubble-detector) -> Core ML. |
| |
| Steps |
| 1. Cut the exported ONNX before its post-processor (TopK / label decode / orig-size scaling) so the |
| Core ML model returns raw per-query outputs: |
| logits [1,300,3] (pre-sigmoid class logits: bubble, text_bubble, text_free) |
| boxes [1,300,4] (cx, cy, w, h normalised to the 640x640 input) |
| 2. Patch Reshape(allowzero=1) -> 0 (no zero-sized dims in this graph) and register the few |
| onnx2torch converters missing for opset 18 (Resize-18, Split-18, ReduceMean/Max-18, GridSample-16). |
| 3. onnx2torch -> torch.jit.trace (fixed 1x3x640x640) -> coremltools mlprogram, fp16 and fp32 |
| (+ optional int8 weight-only palettisation/linear quantisation). |
| 4. Parity vs onnxruntime on the benchmark pages. |
| |
| Usage: python convert_v4s_coreml.py <src.onnx> <out_prefix> |
| e.g. python convert_v4s_coreml.py models/ogkalu-ctbd/detector-v4-s.fp32-5084e7d.onnx coreml/detector-chosen/ogkalu-v4s |
| """ |
| import sys, os, glob |
| import numpy as np, onnx, torch, torch.nn.functional as F |
| from onnx.utils import extract_model |
| from onnx2torch import convert |
| from onnx2torch.node_converters.registry import add_converter, _CONVERTER_REGISTRY, OperationDescription |
| from onnx2torch.utils.common import OperationConverterResult, onnx_mapping_from_node, get_const_value, OnnxToTorchModule |
|
|
| import os |
| B = os.environ.get("NEMU_OCR_WORKDIR", ".") |
| SRC, PREFIX = sys.argv[1], sys.argv[2] |
| LOGITS, BOXES = sys.argv[3] if len(sys.argv) > 3 else "select_36", sys.argv[4] if len(sys.argv) > 4 else "select_37" |
|
|
|
|
| def reg(op, ver, fn): |
| k = OperationDescription(domain="", operation_type=op, version=ver) |
| if k not in _CONVERTER_REGISTRY: add_converter(operation_type=op, version=ver)(fn) |
|
|
|
|
| get = lambda op, v: _CONVERTER_REGISTRY[OperationDescription(domain="", operation_type=op, version=v)] |
| reg("Resize", 18, get("Resize", 13)) |
| reg("Resize", 19, get("Resize", 13)) |
| reg("Split", 18, get("Split", 13)) |
|
|
|
|
| class ReduceStatic(torch.nn.Module, OnnxToTorchModule): |
| def __init__(self, fn, axes, keepdims): |
| super().__init__(); self.fn, self.axes, self.keepdims = fn, axes, bool(keepdims) |
|
|
| def forward(self, x, *_): |
| if self.fn == "mean": return x.mean(dim=self.axes, keepdim=self.keepdims) |
| out = x |
| for a in sorted([a % x.dim() for a in self.axes], reverse=True): out = out.amax(dim=a, keepdim=self.keepdims) |
| return out |
|
|
|
|
| def reduce_conv(fn): |
| def _(node, graph): |
| axes = node.attributes.get("axes") |
| if axes is None and len(node.input_values) > 1: |
| axes = get_const_value(node.input_values[1], graph).tolist() |
| return OperationConverterResult(torch_module=ReduceStatic(fn, list(axes), node.attributes.get("keepdims", 1)), |
| onnx_mapping=onnx_mapping_from_node(node)) |
| return _ |
|
|
|
|
| reg("ReduceMean", 18, reduce_conv("mean")) |
| reg("ReduceMax", 18, reduce_conv("max")) |
|
|
|
|
| class GridSample(torch.nn.Module, OnnxToTorchModule): |
| def __init__(self, mode, padding_mode, align_corners): |
| super().__init__(); self.mode, self.padding_mode, self.align_corners = mode, padding_mode, align_corners |
|
|
| def forward(self, x, grid): |
| return F.grid_sample(x, grid, mode=self.mode, padding_mode=self.padding_mode, align_corners=self.align_corners) |
|
|
|
|
| def gs(node, graph): |
| a = node.attributes |
| mode = {"linear": "bilinear", "bilinear": "bilinear", "nearest": "nearest", "cubic": "bicubic", "bicubic": "bicubic"}[a.get("mode", "linear")] |
| return OperationConverterResult(torch_module=GridSample(mode, a.get("padding_mode", "zeros"), bool(a.get("align_corners", 0))), |
| onnx_mapping=onnx_mapping_from_node(node)) |
|
|
|
|
| reg("GridSample", 16, gs) |
| reg("GridSample", 20, gs) |
|
|
| |
| core = f"{B}/detector/convert/{os.path.basename(PREFIX)}-core.onnx" |
| extract_model(SRC, core, ["images"], [LOGITS, BOXES]) |
| m = onnx.load(core) |
| for n in m.graph.node: |
| if n.op_type == "Reshape": |
| for at in n.attribute: |
| if at.name == "allowzero": at.i = 0 |
| onnx.save(m, core) |
| tm = convert(m).eval() |
|
|
| |
| import onnxruntime as ort |
| from PIL import Image |
| sess = ort.InferenceSession(core, providers=["CPUExecutionProvider"]) |
| def prep(f): |
| return (np.asarray(Image.open(f).convert("RGB").resize((640, 640), Image.BILINEAR), dtype=np.float32).transpose(2, 0, 1)[None] / 255.0) |
| x = prep(f"{B}/pages/hanako-ch1-p02.png") |
| ref = sess.run(None, {"images": x}) |
| with torch.no_grad(): t = tm(torch.from_numpy(x)) |
| print("onnx2torch vs ORT max abs diff:", [float(np.abs(r - tt.numpy()).max()) for r, tt in zip(ref, t)]) |
|
|
|
|
| class Wrap(torch.nn.Module): |
| def __init__(self, m): super().__init__(); self.m = m |
| def forward(self, images): |
| logits, boxes = self.m(images) |
| return logits, boxes |
|
|
|
|
| with torch.no_grad(): |
| traced = torch.jit.trace(Wrap(tm).eval(), torch.from_numpy(x), check_trace=False) |
|
|
| import coremltools as ct |
| for prec in ("fp32", "fp16"): |
| ml = ct.convert(traced, convert_to="mlprogram", |
| inputs=[ct.TensorType(name="images", shape=(1, 3, 640, 640), dtype=np.float32)], |
| outputs=[ct.TensorType(name="logits"), ct.TensorType(name="boxes")], |
| compute_precision=ct.precision.FLOAT16 if prec == "fp16" else ct.precision.FLOAT32, |
| minimum_deployment_target=ct.target.iOS17) |
| ml.short_description = ("ogkalu/comic-text-and-bubble-detector detector-v4-s (RT-DETRv4-S / D-FINE decoder). " |
| "Input: RGB page resized (no letterbox) to 640x640, NCHW float32 in [0,1]. " |
| "Outputs: 300 queries x 3 class logits (bubble, text_bubble, text_free; apply sigmoid) " |
| "and cx,cy,w,h boxes normalised to the page. See postprocess-spec.md.") |
| ml.author = "ogkalu (weights, Apache-2.0); Core ML conversion for nemu" |
| ml.license = "Apache-2.0 (https://huggingface.co/ogkalu/comic-text-and-bubble-detector)" |
| ml.version = "v4-s int8 @ 16e8a622 dequantised" if "deq" in SRC else "v4-s fp32 @ 5084e7d" |
| ml.user_defined_metadata["classes"] = "bubble,text_bubble,text_free" |
| ml.user_defined_metadata["input_size"] = "640x640" |
| ml.user_defined_metadata["source_onnx"] = os.path.basename(SRC) |
| ml.input_description["images"] = "1x3x640x640 float32 RGB in [0,1]; page resized with antialiased bilinear (PIL BILINEAR), no letterbox, no mean/std" |
| ml.output_description["logits"] = "1x300x3 pre-sigmoid logits: bubble, text_bubble, text_free" |
| ml.output_description["boxes"] = "1x300x4 (cx, cy, w, h) normalised to [0,1] of the original page" |
| ml.save(f"{PREFIX}-{prec}.mlpackage") |
| out = ml.predict({"images": x}) |
| print(prec, "Core ML vs ORT max abs diff: logits", float(np.abs(out["logits"] - ref[0]).max()), "boxes", float(np.abs(out["boxes"] - ref[1]).max())) |
|
|
