File size: 3,642 Bytes
3c19d8e | 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 | #!/usr/bin/env python3
"""TableFormerV2 ONNX inference — table structure recognition.
Self-contained: requires only numpy, onnxruntime and Pillow.
python inference.py <table_image.png>
Pipeline: encoder.onnx -> greedy loop over decoder.onnx -> bbox_head.onnx.
The decoder graph is cache-free and re-runs the whole prefix each step
(vocab is 13 tokens, sequences are short, so this is cheap).
Outputs OTSL structure tokens and one bbox per data cell, xyxy normalized
to [0, 1] relative to the input table crop.
"""
import sys
from pathlib import Path
import numpy as np
import onnxruntime as ort
from PIL import Image
HERE = Path(__file__).parent
ID2TOKEN = [
"<pad>", "[UNK]", "<start>", "<end>", "<ecel>", "<fcel>", "<lcel>",
"<ucel>", "<xcel>", "<nl>", "<ched>", "<rhed>", "<srow>",
]
BOS_ID, EOS_ID = 2, 3
DATA_CELL_IDS = {4, 5, 10, 11, 12} # ecel, fcel, ched, rhed, srow
IMAGE_SIZE = 448
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) # ImageNet
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
def preprocess(image: Image.Image) -> np.ndarray:
"""RGB table crop -> (1, 3, 448, 448) float32, ImageNet-normalized."""
image = image.convert("RGB").resize((IMAGE_SIZE, IMAGE_SIZE), Image.BILINEAR)
x = np.asarray(image, dtype=np.float32) / 255.0
x = (x - MEAN) / STD
return x.transpose(2, 0, 1)[None]
class TableFormerV2Onnx:
def __init__(self, model_dir=HERE, providers=("CPUExecutionProvider",)):
model_dir = Path(model_dir)
p = list(providers)
self.encoder = ort.InferenceSession(str(model_dir / "encoder.onnx"), providers=p)
self.decoder = ort.InferenceSession(str(model_dir / "decoder.onnx"), providers=p)
self.bbox_head = ort.InferenceSession(str(model_dir / "bbox_head.onnx"), providers=p)
def predict(self, image: Image.Image, max_length: int = 512):
"""Returns (otsl_tokens, bboxes). bboxes: (num_cells, 4) xyxy in [0, 1]."""
images = preprocess(image)
(enc_hidden,) = self.encoder.run(None, {"images": images})
# Greedy generation
ids = np.array([[BOS_ID]], dtype=np.int64)
for _ in range(max_length):
logits, _ = self.decoder.run(
None, {"input_ids": ids, "encoder_hidden": enc_hidden}
)
next_id = int(logits[0, -1].argmax())
ids = np.concatenate([ids, [[next_id]]], axis=1)
if next_id == EOS_ID:
break
# Hidden states of the full sequence -> bboxes at data-cell positions
_, hidden = self.decoder.run(
None, {"input_ids": ids, "encoder_hidden": enc_hidden}
)
cell_pos = [i for i, t in enumerate(ids[0].tolist()) if t in DATA_CELL_IDS]
if cell_pos:
(bboxes,) = self.bbox_head.run(
None,
{"cell_embeddings": hidden[0, cell_pos], "encoder_hidden": enc_hidden},
)
else:
bboxes = np.zeros((0, 4), dtype=np.float32)
tokens = [
ID2TOKEN[t] for t in ids[0].tolist() if t not in (BOS_ID, EOS_ID, 0)
]
return tokens, bboxes
if __name__ == "__main__":
if len(sys.argv) != 2:
sys.exit(f"usage: {sys.argv[0]} <table_image>")
img = Image.open(sys.argv[1])
model = TableFormerV2Onnx()
tokens, bboxes = model.predict(img)
print("OTSL:", " ".join(tokens))
print(f"{len(bboxes)} cells (xyxy in original image pixels):")
scale = np.array([img.width, img.height, img.width, img.height])
for b in bboxes * scale:
print(f" ({b[0]:7.1f}, {b[1]:7.1f}, {b[2]:7.1f}, {b[3]:7.1f})")
|