| """Run AE3-TL on a Mac with Core ML, and compare it with the ONNX CPU path. |
| |
| python predict_mac.py page.jpg |
| python predict_mac.py page.jpg --units all |
| python predict_mac.py page.jpg --compare-onnx |
| |
| The page model runs in Core ML; the relation head stays on the CPU with ONNX Runtime. Table boxes and |
| caption and footer links are printed. Full cells and text come from the kit's post-processing. |
| """ |
| import argparse |
| import time |
|
|
| import numpy as np |
| from PIL import Image |
| import onnxruntime as ort |
|
|
| try: |
| import coremltools as ct |
| except ImportError: |
| raise SystemExit("coremltools is required: pip install 'coremltools>=8'") |
|
|
| PACKAGE = "coreml/ae3tl_pages_v3_ep1_ema_fp32.mlpackage" |
| REL = "weights/ae3tl_relhead_v3_ep1_fp32_opset17.onnx" |
| PAGES_ONNX = "weights/ae3tl_pages_v3_ep1_fp32_opset17.onnx" |
| TABLE, CAPTION, FOOTER = (0, 8), (6, 14), (7, 15) |
| UNITS = { |
| "cpu": ct.ComputeUnit.CPU_ONLY, |
| "gpu": ct.ComputeUnit.CPU_AND_GPU, |
| "all": ct.ComputeUnit.ALL, |
| "ane": ct.ComputeUnit.CPU_AND_NE, |
| } |
|
|
|
|
| def preprocess(path): |
| img = Image.open(path).convert("RGB") |
| W, H = img.size |
| x = np.asarray(img.resize((768, 768), Image.BILINEAR), dtype=np.float32) / 255.0 |
| return W, H, x.transpose(2, 0, 1)[None] |
|
|
|
|
| def decode(logits, boxes, W, H, threshold=0.5): |
| p = 1.0 / (1.0 + np.exp(-logits[0])) |
| order = np.argsort(-p.reshape(-1))[:300] |
| scores = p.reshape(-1)[order] |
| labels, qids = (order % 16), (order // 16) |
| b = boxes[0][qids] |
| xyxy = np.stack([b[:, 0] - b[:, 2] / 2, b[:, 1] - b[:, 3] / 2, |
| b[:, 0] + b[:, 2] / 2, b[:, 1] + b[:, 3] / 2], 1) |
| xyxy = xyxy * np.array([W, H, W, H], dtype=np.float32) |
| keep = scores >= threshold |
| labels, qids, xyxy = labels[keep], qids[keep], xyxy[keep] |
| tables = [(int(q), box.tolist()) for q, box in zip(qids[np.isin(labels, TABLE)], xyxy[np.isin(labels, TABLE)])] |
| captions = [int(q) for q in qids[np.isin(labels, CAPTION)]] |
| footers = [int(q) for q in qids[np.isin(labels, FOOTER)]] |
| return tables, captions, footers |
|
|
|
|
| def run_coreml(path, units): |
| model = ct.models.MLModel(PACKAGE, compute_units=UNITS[units]) |
| W, H, x = preprocess(path) |
| t0 = time.perf_counter() |
| out = model.predict({"images": x}) |
| dt = time.perf_counter() - t0 |
| tables, captions, footers = decode(out["logits"], out["boxes"], W, H) |
| return tables, captions, footers, out["query_embeddings"], dt |
|
|
|
|
| def run_onnx(path): |
| session = ort.InferenceSession(PAGES_ONNX, providers=["CPUExecutionProvider"]) |
| W, H, x = preprocess(path) |
| t0 = time.perf_counter() |
| out = dict(zip([o.name for o in session.get_outputs()], session.run(None, {"images": x}))) |
| dt = time.perf_counter() - t0 |
| tables, captions, footers = decode(out["logits"], out["boxes"], W, H) |
| return tables, captions, footers, out["query_embeddings"], dt |
|
|
|
|
| def link(embeddings, tables, quarters, rel): |
| """The relation head on the CPU: caption or footer query against each table query.""" |
| result = {} |
| for q in quarters: |
| pairs = [] |
| for tq, _ in tables: |
| if tq != q: |
| pairs += [[int(q), int(tq)], [int(tq), int(q)]] |
| if not pairs: |
| result[q] = None |
| continue |
| e = rel.run(None, {"query_embeddings": embeddings[0], |
| "pairs": np.asarray(pairs, dtype=np.int64)})[0] |
| e = 1.0 / (1.0 + np.exp(-e)) |
| best, best_score = None, -1.0 |
| for k, (tq, _) in enumerate(tables): |
| m = float((e[2 * k] + e[2 * k + 1]) / 2.0) |
| if m > best_score: |
| best, best_score = tq, m |
| result[q] = (best, round(best_score, 3)) |
| return result |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("page") |
| ap.add_argument("--units", default="gpu", choices=sorted(UNITS)) |
| ap.add_argument("--compare-onnx", action="store_true") |
| args = ap.parse_args() |
|
|
| rel = ort.InferenceSession(REL, providers=["CPUExecutionProvider"]) |
| tables, captions, footers, emb, dt = run_coreml(args.page, args.units) |
| print(f"compute unit: {args.units}") |
| print(f"page model: {dt * 1000:.1f} ms") |
| print(f"tables: {len(tables)}") |
| for q, box in tables: |
| print(" table", q, [round(v) for v in box]) |
| for kind, quarters in (("caption", captions), ("footer", footers)): |
| for q, (best, score) in link(emb, tables, quarters, rel).items(): |
| print(f" {kind} {q} -> table {best} score {score}") |
|
|
| if args.compare_onnx: |
| t2, c2, f2, emb2, dt2 = run_onnx(args.page) |
| print(f"\nONNX CPU: {dt2 * 1000:.1f} ms, tables {len(t2)}") |
| same_tables = len(tables) == len(t2) |
| max_dev = 0.0 |
| for (q1, b1), (q2, b2) in zip(tables, t2): |
| max_dev = max(max_dev, max(abs(a - b) for a, b in zip(b1, b2))) |
| print(f"same table count: {same_tables}; max box difference: {max_dev:.1f} px") |
| l1 = link(emb, tables, captions, rel) |
| l2 = link(emb2, t2, c2, rel) |
| print(f"same caption links: {l1 == l2}; same footer links: {link(emb, tables, footers, rel) == link(emb2, t2, f2, rel)}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
