File size: 4,521 Bytes
f8b3943 eb6c45c f8b3943 d02201e f8b3943 d02201e ef627b3 9f2d61d ef627b3 9f2d61d d02201e 9f2d61d ef627b3 d02201e ef627b3 d02201e f8b3943 d02201e f8b3943 | 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 | """Minimal AE3-TL example: page model and relation head with ONNX Runtime.
Prints the table boxes and, for each caption and footer box, the table it links to.
For full tables with cells and text, run extract.py, which uses src/ae3_tl.
"""
import os
import sys
import numpy as np
from PIL import Image
import onnxruntime as ort
PAGES = "weights/ae3tl_pages_v3_ep1_fp32_opset17.onnx"
REL = "weights/ae3tl_relhead_v3_ep1_fp32_opset17.onnx"
TABLE, CAPTION, FOOTER = (0, 8), (6, 14), (7, 15)
def _provider_list():
"""Pick the fastest provider this machine has. AE3_PROVIDER overrides the choice."""
forced = os.environ.get("AE3_PROVIDER", "auto")
have = ort.get_available_providers()
if forced in ("cpu", "cuda", "coreml"):
choice = forced
elif "CUDAExecutionProvider" in have:
choice = "cuda"
elif "CoreMLExecutionProvider" in have:
choice = "coreml"
else:
choice = "cpu"
print(f"provider: {choice} (available: {', '.join(have)})")
if choice == "coreml":
cache = os.environ.get("AE3_COREML_CACHE", os.path.expanduser("~/.cache/ae3tl-coreml"))
return [("CoreMLExecutionProvider", {"ModelFormat": "MLProgram", "MLComputeUnits": "ALL",
"RequireStaticInputShapes": "1", "ModelCacheDirectory": cache}),
"CPUExecutionProvider"]
if choice == "cuda":
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
return ["CPUExecutionProvider"]
def _session(path, prov):
"""Create a session. If CoreML is chosen it must start, or we stop with the reason."""
coreml = any("CoreML" in str(p) for p in prov)
model, opts = path, None
if coreml:
if "CoreMLExecutionProvider" not in ort.get_available_providers():
sys.exit("CoreMLExecutionProvider is missing from this onnxruntime build. On Apple "
"Silicon install the standard wheel: pip install -U 'onnxruntime>=1.20'.")
b1 = os.path.splitext(path)[0] + "_b1.onnx"
if os.path.exists(b1):
model = b1
else:
opts = ort.SessionOptions()
opts.add_free_dimension_override_by_name("batch", 1)
try:
return (ort.InferenceSession(model, opts, providers=prov) if opts
else ort.InferenceSession(model, providers=prov))
except Exception as exc:
if not coreml:
raise
sys.exit(f"CoreML failed to start: {type(exc).__name__}: {exc}")
def main(path):
prov = _provider_list()
pages = _session(PAGES, prov)
rel = _session(REL, prov)
img = Image.open(path).convert("RGB")
W, H = img.size
x = np.asarray(img.resize((768, 768), Image.BILINEAR), dtype=np.float32) / 255.0
x = x.transpose(2, 0, 1)[None]
logits, boxes, emb, _ = pages.run(None, {"images": x})
p = 1.0 / (1.0 + np.exp(-logits[0])) # queries x 16 classes
order = np.argsort(-p.reshape(-1))[:300]
scores = p.reshape(-1)[order]
labels, qids = (order % 16), (order // 16)
b = boxes[0][qids] # cx, cy, w, h, normalised
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 >= 0.5
labels, qids, xyxy = labels[keep], qids[keep], xyxy[keep]
tab = np.isin(labels, TABLE)
tables = [(int(q), box.tolist()) for q, box in zip(qids[tab], xyxy[tab])]
print("tables:", len(tables))
for q, box in tables:
print(" table", q, [round(v) for v in box])
for kind, classes in (("caption", CAPTION), ("footer", FOOTER)):
sel = np.isin(labels, classes)
for q in qids[sel]:
pairs = []
for tq, _ in tables:
if tq != q:
pairs += [[int(q), int(tq)], [int(tq), int(q)]]
if not pairs:
print(kind, q, "no tables"); continue
e = rel.run(None, {"query_embeddings": emb[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
print(kind, q, "-> table", best, "score", round(best_score, 3))
if __name__ == "__main__":
main(sys.argv[1])
|