| """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])) |
| 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 >= 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]) |
|
|
