File size: 5,157 Bytes
896f580 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """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()
|