tableformer_accurate/inference.py

#!/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})")