File size: 6,610 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 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | #!/usr/bin/env python3
"""TableFormer v1 ONNX inference — table structure recognition.
Self-contained: requires only numpy, onnxruntime and Pillow.
The same script works for both the fast and accurate variants
(the decoder layer count is read from the cache input shape).
python inference.py <table_image.png>
Pipeline: encoder.onnx -> greedy loop over decoder_step.onnx -> bbox_decoder.onnx.
The decoder computes one token per call; `cache` carries each decoder layer's
per-position outputs across steps (shape (num_layers, L, 1, 512); pass length 0
on the first step).
Replicates TableModel04_rs.predict from docling-ibm-models, including its
structure-error corrections and horizontal-span bbox merging.
Outputs OTSL structure tokens and one bbox per cell, cxcywh 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
WM = {
"<pad>": 0, "<unk>": 1, "<start>": 2, "<end>": 3, "ecel": 4, "fcel": 5,
"lcel": 6, "ucel": 7, "xcel": 8, "nl": 9, "ched": 10, "rhed": 11, "srow": 12,
}
ID2TOKEN = {v: k for k, v in WM.items()}
MAX_STEPS = 1024
IMAGE_SIZE = 448 # resized without keeping aspect ratio (per tm_config.json)
MEAN = np.array([0.94247851, 0.94254675, 0.94292611], dtype=np.float32)
STD = np.array([0.17910956, 0.17940403, 0.17931663], dtype=np.float32)
def preprocess(image: Image.Image) -> np.ndarray:
"""RGB table crop -> (1, 3, 448, 448) float32, PubTabNet-normalized.
Note: docling's TFPredictor._prepare_image feeds the image TRANSPOSED,
as (channels, width, height); the model was trained that way. Predicted
bbox x/y are therefore swapped back in predict().
"""
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, 1, 0)[None] # (1, C, W, H)
class TableFormerOnnx:
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_step.onnx"), providers=p)
self.bbox = ort.InferenceSession(str(model_dir / "bbox_decoder.onnx"), providers=p)
cache_shape = next(
i.shape for i in self.decoder.get_inputs() if i.name == "cache"
)
self.num_layers, self.hidden_dim = cache_shape[0], cache_shape[3]
def predict(self, image: Image.Image):
"""Returns (otsl_tokens, classes, bboxes).
classes: (num_cells, 3) logits; bboxes: (num_cells, 4) cxcywh in [0, 1].
"""
x = preprocess(image)
enc_out, memory = self.encoder.run(None, {"image": x})
decoded = np.array([[WM["<start>"]]], dtype=np.int64) # (L, 1)
cache = np.zeros((self.num_layers, 0, 1, self.hidden_dim), dtype=np.float32)
output_tags, tag_H_buf = [], []
skip_next_tag, prev_tag_ucel = True, False
first_lcel, bboxes_to_merge, cur_bbox_ind, bbox_ind = True, {}, -1, 0
line_num = 0
while len(output_tags) < MAX_STEPS:
logits, tag_H, cache = self.decoder.run(
None, {"decoded_tags": decoded, "memory": memory, "cache": cache}
)
new_tag = int(logits[0].argmax())
# Structure error corrections (as in TableModel04_rs.predict)
if line_num == 0 and new_tag == WM["xcel"]:
new_tag = WM["lcel"]
if prev_tag_ucel and new_tag == WM["lcel"]:
new_tag = WM["fcel"]
if new_tag == WM["<end>"]:
output_tags.append(new_tag)
break
output_tags.append(new_tag)
if new_tag == WM["nl"]:
line_num += 1
# Keep one hidden state per cell for the bbox decoder
if not skip_next_tag and new_tag in (
WM["fcel"], WM["ecel"], WM["ched"], WM["rhed"],
WM["srow"], WM["nl"], WM["ucel"],
):
tag_H_buf.append(tag_H)
if not first_lcel:
bboxes_to_merge[cur_bbox_ind] = bbox_ind
bbox_ind += 1
if new_tag != WM["lcel"]:
first_lcel = True
elif first_lcel: # start of a horizontal span
tag_H_buf.append(tag_H)
first_lcel = False
cur_bbox_ind = bbox_ind
bboxes_to_merge[cur_bbox_ind] = -1
bbox_ind += 1
skip_next_tag = new_tag in (WM["nl"], WM["ucel"], WM["xcel"])
prev_tag_ucel = new_tag == WM["ucel"]
decoded = np.concatenate([decoded, [[new_tag]]], axis=0)
tokens = [ID2TOKEN[t] for t in output_tags if t != WM["<end>"]]
if not tag_H_buf:
return tokens, np.zeros((0, 3), np.float32), np.zeros((0, 4), np.float32)
classes, coords = self.bbox.run(
None, {"enc_out": enc_out, "tag_H": np.concatenate(tag_H_buf, axis=0)}
)
# Merge first/last bbox of each horizontal span (cxcywh)
out_cls, out_coord, skip = [], [], set()
for i in range(len(coords)):
if i in bboxes_to_merge:
b1, b2 = coords[i], coords[bboxes_to_merge[i]]
skip.add(bboxes_to_merge[i])
w = (b2[0] + b2[2] / 2) - (b1[0] - b1[2] / 2)
h = (b2[1] + b2[3] / 2) - (b1[1] - b1[3] / 2)
left = b1[0] - b1[2] / 2
top = min(b2[1] - b2[3] / 2, b1[1] - b1[3] / 2)
out_coord.append([left + w / 2, top + h / 2, w, h])
out_cls.append(classes[i])
elif i not in skip:
out_coord.append(coords[i].tolist())
out_cls.append(classes[i])
return tokens, np.array(out_cls), np.array(out_coord, dtype=np.float32)
if __name__ == "__main__":
if len(sys.argv) != 2:
sys.exit(f"usage: {sys.argv[0]} <table_image>")
img = Image.open(sys.argv[1])
model = TableFormerOnnx()
tokens, classes, bboxes = model.predict(img)
print("OTSL:", " ".join(tokens))
print(f"{len(bboxes)} cells (xyxy in original image pixels):")
for cx, cy, w, h in bboxes:
x1, y1 = (cx - w / 2) * img.width, (cy - h / 2) * img.height
x2, y2 = (cx + w / 2) * img.width, (cy + h / 2) * img.height
print(f" ({x1:7.1f}, {y1:7.1f}, {x2:7.1f}, {y2:7.1f})")
|