Simpler model card with benchmark images; one line per recognizer call on CPU (415 ms/page)

8941e3a verified 4 days ago

20.7 kB

"""End-to-end document processing: file -> pages -> text segments -> labelled fields -> JSON."""

import json
import time
from pathlib import Path

import cv2
import numpy as np

from ..charset import NUM_CLASSES, ctc_greedy_decode
from ..data.datasets import line_to_input, orientation_page, text_patches
from ..features import MAX_SEGMENTS, segment_features
from ..geometry import crop_quad, expand_quad, quad_bbox, quad_size, transform_quads
from ..models.orientation import ANGLES, PATCH
from ..models.recognizer import WIDTH_STRIDE
from ..schema import DOC_TYPES, FIELD_LABELS, REL_TYPES
from .assemble import (assemble_fields, assemble_page_structure, block_groups, block_order, groups_from_block_map,
                       structure_from_groups, xy_cut)
from .detect import DET_MAX_SIDE, detector_input, find_rules, kernels_to_quads, split_gutter_crossers
from .inputs import estimate_skew, load_pages, rot90_matrix, rotate_image

COMPONENTS = ("detector", "recognizer", "orientation", "extractor")

# Frequent words used to check a proposed page rotation by reading a few lines both ways (upside-down text
# still yields plausible-looking letters, so OCR confidence alone does not separate the two).
COMMON_WORDS = frozenset("""
the of and to in is for on that with as by this be are from at or an it was will not have all can has
but which their more one were been other its into also than about after these such may when only over
new they we you he she his her our your there two first time most some would what so if no up out who
total tax subtotal cash change qty amount price date invoice item items paid due balance receipt number
page table figure section chapter results data used use between each both during through under where
""".split())
SCHEMA_VERSION = "1.1"   # 1.1: "formula" blocks; layout blocks from the model's same_block relation
# White columns after every line the recognizer reads. Training pads each line of a batch with white to the right
# (to a multiple of 64 px), so the model reads a line's end against white; without this strip a line read on its
# own, or the widest line of a batch, would end at the image edge instead. 32 px covers the CNN's receptive field
# (64 px gives identical output), and with it the output no longer depends on how lines are batched.
RECOGNIZER_PAD = 32


def _softmax(x, axis=-1):
    e = np.exp(x - x.max(axis, keepdims=True))
    return e / e.sum(axis, keepdims=True)


def line_batch(lines):
    """Recognizer input for line arrays (32, W) in [0, 1]: (batch padded with white, step mask, steps per line)."""
    wmax = max(x.shape[1] for x in lines) + RECOGNIZER_PAD
    batch = np.ones((len(lines), 1, lines[0].shape[0], wmax), np.float32)
    steps = np.array([x.shape[1] // WIDTH_STRIDE for x in lines])
    for k, x in enumerate(lines):
        batch[k, 0, :, : x.shape[1]] = x
    mask = np.arange(wmax // WIDTH_STRIDE)[None, :] >= steps[:, None]
    return batch, mask, steps


def _sigmoid(x):
    return 1 / (1 + np.exp(-np.clip(x, -60, 60)))


class TorchBackend:
    def __init__(self, weights_dir, device="cpu"):
        import torch

        from ..models import build, count_params

        self.torch, self.device, self.models, self.params = torch, torch.device(device), {}, {}
        for name in COMPONENTS:
            p = Path(weights_dir) / name / "best.pt"
            if not p.exists():
                raise FileNotFoundError(f"missing {p}; train it with `python -m tinydoc.train {name}`")
            ck = torch.load(p, map_location="cpu", weights_only=False)
            m = build(name, **(ck.get("arch") or {}))
            m.load_state_dict(ck["model"])
            self.models[name] = m.to(self.device).eval()
            self.params[name] = count_params(m)
        self.vocab = NUM_CLASSES

    def __call__(self, name, *inputs):
        with self.torch.no_grad():
            t = [self.torch.from_numpy(np.ascontiguousarray(x)).to(self.device) for x in inputs]
            out = self.models[name](*t)
            out = out if isinstance(out, tuple) else (out,)
            return [o.float().cpu().numpy() for o in out]


class OnnxBackend:
    def __init__(self, onnx_dir, threads=None):
        import onnx
        import onnxruntime as ort

        opts = ort.SessionOptions()
        if threads:
            opts.intra_op_num_threads = threads
        self.sessions, self.params = {}, {}
        for name in COMPONENTS:
            p = Path(onnx_dir) / f"{name}.onnx"
            if not p.exists():
                raise FileNotFoundError(f"missing {p}; export with `python -m tinydoc.export`")
            self.sessions[name] = ort.InferenceSession(str(p), opts, providers=["CPUExecutionProvider"])
            meta = {m.key: m.value for m in onnx.load(str(p), load_external_data=False).metadata_props}
            self.params[name] = int(meta.get("trained_parameters", 0))
        # Character vocabulary the models were trained with: releases before the inline-maths characters were
        # appended have fewer classes, and their extractor must not see the newer ids.
        width = self.sessions["recognizer"].get_outputs()[0].shape[-1]
        self.vocab = width if isinstance(width, int) else NUM_CLASSES

    def __call__(self, name, *inputs):
        s = self.sessions[name]
        feed = {i.name: np.ascontiguousarray(x) for i, x in zip(s.get_inputs(), inputs)}
        return s.run(None, feed)


class Engine:
    def __init__(self, weights="runs", backend="torch", device="cpu", recognizer_batch=None, model_blocks=None, blocks=None):
        self.backend = TorchBackend(weights, device) if backend == "torch" else OnnxBackend(weights)
        # Lines per recognizer call. On CPU (ONNX Runtime) one line per call is fastest: nothing is padded to the
        # widest line of a batch and the activations stay in cache. On 40 OmniDocBench pages (Apple M5 Pro) it cut
        # recognition from 421 to 248 ms per page at 4 threads and from 1190 to 440 ms at 1 thread, against
        # batches of 32, with the same output (see RECOGNIZER_PAD). The torch backend keeps batches, which pay
        # off on a GPU.
        self.recognizer_batch = recognizer_batch or (1 if backend == "onnx" else 32)
        # How lines are grouped into paragraphs/blocks on ordinary pages: "geometric" rules, the extractor's
        # same_block relation ("model"), or the detector's block regions ("detector"). The default comes from
        # config.json next to the weights, chosen by measuring each on validation pages.
        cfg = Path(weights) / "config.json"
        cfg = json.loads(cfg.read_text()) if cfg.exists() else {}
        if blocks is None and model_blocks is not None:
            blocks = "model" if model_blocks else "geometric"
        self.blocks = blocks or cfg.get("blocks") or ("model" if cfg.get("model_blocks") else "geometric")
        # Splitting lines merged across a column gutter: off by default (on OmniDocBench such merges were rare,
        # 5 in 380 pages, and the split hurt as often as it helped).
        self.split_gutters = cfg.get("split_gutters", False)
        # Detector input size: one pass with the long side at det_max_side. Optionally (det_fast_side) a page is
        # first detected at a smaller size and again at det_max_side only when its median line is under
        # det_min_line_px there. Off by default: on v3 it was ~15% faster and helped academic and exam pages, but
        # the official OmniDocBench evaluator showed it losing small print on newspapers (+0.017 to +0.038 text
        # edit) and magazines (+0.02), which the median-line test does not catch.
        self.det_max_side = cfg.get("det_max_side", DET_MAX_SIDE)
        self.det_fast_side = cfg.get("det_fast_side")
        self.det_min_line_px = cfg.get("det_min_line_px", 14)

    # ------------------------------------------------------------------------------------ stages
    def orientation(self, img):
        gray = orientation_page(img)
        pos = text_patches(gray, 8)
        patches = np.stack([gray[y:y + PATCH, x:x + PATCH] for x, y in pos]).astype(np.float32)[:, None] / 255
        if patches.shape[-1] != PATCH or patches.shape[-2] != PATCH:
            return 0, 1.0
        probs = _softmax(self.backend("orientation", patches)[0]).mean(0)
        k = int(probs.argmax())
        return k, float(probs[k])

    def _reading_evidence(self, img, n=16):
        """(common-word hits, mean OCR confidence) over the n widest detected lines of an image."""
        dets = sorted(self.detect(img), key=lambda d: -quad_size(d[0])[0])[:n]
        if not dets:
            return 0, 0.0
        texts = self.recognize(img, [q for q, _ in dets])
        words = [w.strip(".,:;!?()'\"").lower() for t, _ in texts for w in t.split()]
        return sum(w in COMMON_WORDS for w in words), float(np.mean([c for _, c in texts]))

    def verify_rotation(self, img, k):
        """Keep a proposed 90/180/270 correction only if the rotated page reads better than the original."""
        hits0, conf0 = self._reading_evidence(img)
        hitsk, confk = self._reading_evidence(np.ascontiguousarray(np.rot90(img, k)))
        if abs(hitsk - hits0) >= 2:
            return hitsk > hits0
        if abs(confk - conf0) > 0.05:
            return confk > conf0
        return True   # no clear evidence either way: trust the orientation classifier

    def detect_full(self, img, blocks=True):
        """Line detections plus the block-region map (None for detectors without a block channel, or when
        blocks=False because the caller does not use it)."""
        for side in ([self.det_fast_side] if self.det_fast_side else []) + [self.det_max_side]:
            x, scale = detector_input(img, max_side=side)
            out = self.backend("detector", x[None])[0][0]
            prob = _sigmoid(out[0])
            quads = kernels_to_quads(prob)
            if side == self.det_max_side or not quads or np.median([quad_size(q)[1] for q, _ in quads]) >= self.det_min_line_px:
                break
        block = _sigmoid(out[1]) if blocks and out.shape[0] > 1 else None
        return [(q / scale, s) for q, s in quads], block, scale

    def detect(self, img):
        return self.detect_full(img)[0]

    def recognize(self, img, quads):
        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        crops = []
        for q in quads:
            _, h = quad_size(q)
            crop = crop_quad(gray, expand_quad(q, 0.25 * h, 0.1 * h, 0.25 * h, 0.1 * h))
            if np.median(crop[::2, ::2]) < 100:  # light text on a dark band
                crop = 255 - crop
            crops.append(line_to_input(crop))
        results = [None] * len(crops)
        order = sorted(range(len(crops)), key=lambda i: crops[i].shape[1])
        for start in range(0, len(order), self.recognizer_batch):
            idx = order[start:start + self.recognizer_batch]
            batch, mask, steps = line_batch([crops[i] for i in idx])
            probs = _softmax(self.backend("recognizer", batch, mask)[0])
            for k, i in enumerate(idx):
                text, confs = ctc_greedy_decode(probs[k, : steps[k]])
                results[i] = (text, float(np.mean(confs)) if confs else 0.0)
        return results

    def extract(self, segs, W, H):
        if not segs:
            return np.zeros((0, len(FIELD_LABELS))), np.zeros((len(REL_TYPES), 0, 0)), np.full(len(DOC_TYPES), 1 / len(DOC_TYPES))
        n = len(segs)
        labels = np.zeros((n, len(FIELD_LABELS)), np.float32)
        rel = None   # sized from the model: older exports predict fewer relations
        docs = []
        # windows of MAX_SEGMENTS in reading order, overlapping by half, for very dense pages
        step = MAX_SEGMENTS // 2
        starts = list(range(0, max(n - MAX_SEGMENTS, 0) + 1, step)) or [0]
        if starts[-1] + MAX_SEGMENTS < n:
            starts.append(n - MAX_SEGMENTS)
        for s0 in starts:
            idx = list(range(s0, min(n, s0 + MAX_SEGMENTS)))
            f = segment_features([segs[i] for i in idx], W, H)
            f["char_ids"][f["char_ids"] >= getattr(self.backend, "vocab", NUM_CLASSES)] = 0
            out = self.backend("extractor", f["char_ids"][None], f["feats"][None], f["rel_boxes"][None],
                               np.ones((1, len(idx)), np.float32))
            lp = _softmax(out[0][0])
            labels[idx, :lp.shape[1]] = np.maximum(labels[idx, :lp.shape[1]], lp)
            if rel is None:
                rel = np.zeros((out[1].shape[1], n, n), np.float32)
            rel[:, idx[0]:idx[-1] + 1, idx[0]:idx[-1] + 1] = np.maximum(rel[:, idx[0]:idx[-1] + 1, idx[0]:idx[-1] + 1], _sigmoid(out[1][0]))
            docs.append(_softmax(out[2][0]))
        return labels, rel, np.mean(docs, 0)

    # -------------------------------------------------------------------------------------- run
    def process(self, path, force_ocr=False):
        t_start = time.time()
        timings = {"load": 0.0, "orientation": 0.0, "deskew": 0.0, "detect": 0.0, "recognize": 0.0, "extract": 0.0, "assemble": 0.0}
        pages_out, all_segs, page_info = [], [], []
        t = time.time()
        pages = list(load_pages(path, force_ocr=force_ocr))
        timings["load"] += time.time() - t

        for page in pages:
            info = {"page": page.number, "width": page.width, "height": page.height, "text_source": page.text_source}
            if page.text_source == "pdf_text":
                segs = [{"bbox": [x0, y0, x1, y1], "quad": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], "text": t_, "conf": 1.0}
                        for x0, y0, x1, y1, t_ in page.pdf_segments]
                frame_w, frame_h = page.width, page.height
                rules, block_map, det_scale = [], None, 1.0
            else:
                img = page.image
                t = time.time()
                k, kconf = self.orientation(img)
                if k != 0 and kconf > 0.6 and not self.verify_rotation(img, k):
                    k, kconf = 0, 1.0 - kconf
                    info["rotation_rejected"] = True
                timings["orientation"] += time.time() - t
                m_total = np.eye(3)
                if k != 0 and kconf > 0.6:
                    m_total = np.vstack([rot90_matrix(k, img.shape[1], img.shape[0]), [0, 0, 1]]) @ m_total
                    img = np.ascontiguousarray(np.rot90(img, k))
                info["rotation_corrected_deg"] = ANGLES[k] if (k != 0 and kconf > 0.6) else 0
                info["orientation_confidence"] = round(kconf, 3)
                t = time.time()
                skew = estimate_skew(img)  # text is rotated `skew` degrees counter-clockwise
                if abs(skew) >= 0.5:
                    img, m = rotate_image(img, -skew)  # rotate it back
                    m_total = np.vstack([m, [0, 0, 1]]) @ m_total
                info["skew_corrected_deg"] = round(skew, 2) if abs(skew) >= 0.5 else 0.0
                timings["deskew"] += time.time() - t
                t = time.time()
                dets, block_map, det_scale = self.detect_full(img, blocks=self.blocks == "detector")
                if self.split_gutters:
                    dets = split_gutter_crossers(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), dets)
                timings["detect"] += time.time() - t
                t = time.time()
                texts = self.recognize(img, [q for q, _ in dets])
                timings["recognize"] += time.time() - t
                inv = np.linalg.inv(m_total)
                segs = []
                for (q, dscore), (text, conf) in zip(dets, texts):
                    if not text.strip():
                        continue
                    orig_q = transform_quads([q], inv)[0]   # report positions in the source image
                    segs.append({"bbox": quad_bbox(q), "quad_src": orig_q, "text": text.strip(),
                                 "conf": conf * min(1.0, dscore / 0.7)})
                frame_w, frame_h = img.shape[1], img.shape[0]
                hs = [s["bbox"][3] - s["bbox"][1] for s in segs]
                rules = find_rules(img, float(np.median(hs))) if hs else []
            for i, s in enumerate(segs):
                s["page"] = page.number
                s["id"] = f"p{page.number}-s{i}"
            t = time.time()
            labels, rel, doc = self.extract(segs, frame_w, frame_h)
            timings["extract"] += time.time() - t
            boxes = [s["bbox"] for s in segs]
            groups = None
            general = bool(segs) and DOC_TYPES[int(doc.argmax())] == "general"
            if general and self.blocks == "detector" and block_map is not None:
                groups = block_order(boxes, groups_from_block_map(boxes, block_map, det_scale), rules)
            elif general and self.blocks == "model" and rel.shape[0] > REL_TYPES.index("same_block"):
                names = [FIELD_LABELS[int(r.argmax())] for r in labels]
                groups = block_order(boxes, block_groups(boxes, names, rel[REL_TYPES.index("same_block")]), rules)
            if groups is not None:
                order = [i for g in groups for i in g]
            else:
                order = xy_cut(boxes, rules=rules) if segs else []
            page_info.append((info, segs, labels, rel, doc, order, groups))

        t = time.time()
        doc_probs = np.mean([p[4] for p in page_info if len(p[1])], 0) if any(len(p[1]) for p in page_info) else np.full(len(DOC_TYPES), 1 / len(DOC_TYPES))
        doc_type = DOC_TYPES[int(doc_probs.argmax())]

        # global segment list across pages for field assembly
        offset = 0
        glob_segs, glob_labels, glob_conf, glob_order = [], [], [], []
        n_total = sum(len(p[1]) for p in page_info)
        same_row = np.zeros((n_total, n_total), np.float32)
        key_value = np.zeros((n_total, n_total), np.float32)
        for info, segs, labels, rel, _, order, groups in page_info:
            n = len(segs)
            glob_segs += segs
            glob_labels += [FIELD_LABELS[int(r.argmax())] for r in labels]
            glob_conf += [float(r.max()) for r in labels]
            glob_order += [offset + i for i in order]
            if n:
                same_row[offset:offset + n, offset:offset + n] = rel[0]
                key_value[offset:offset + n, offset:offset + n] = rel[1]
            label_names = [FIELD_LABELS[int(r.argmax())] for r in labels]
            if groups is not None:
                blocks, tables = structure_from_groups(segs, groups, label_names, rel[0])
            else:
                blocks, tables = assemble_page_structure(segs, order, label_names, rel[0] if n else np.zeros((0, 0)))
            pages_out.append({**info, "blocks": blocks, "tables": tables, "segments": [
                {"id": s["id"], "text": s["text"], "bbox": [round(v, 1) for v in (quad_bbox(s["quad_src"]) if "quad_src" in s else s["bbox"])],
                 "quad": np.round(np.asarray(s.get("quad_src", s.get("quad"))), 1).tolist(),
                 "ocr_confidence": round(s["conf"], 3), "label": label_names[i], "label_confidence": round(float(labels[i].max()), 3)}
                for i, s in enumerate(segs)]})
            offset += n

        result = {
            "schema_version": SCHEMA_VERSION,
            "source": {"file": str(Path(path).name), "pages": len(pages)},
            "document_type": {"value": doc_type, "confidence": round(float(doc_probs.max()), 3),
                              "scores": {d: round(float(p), 3) for d, p in zip(DOC_TYPES, doc_probs)}},
        }
        if n_total == 0:
            result["document_type"] = {"value": None, "confidence": 0.0, "scores": None}
            result["fields"] = None
            result["validation"] = {"checks": [], "issues": ["no text found"], "needs_review": True}
        elif doc_type in ("receipt", "invoice"):
            fields, validation = assemble_fields(glob_segs, glob_order, glob_labels, glob_conf, same_row, key_value)
            result["fields"] = fields
            result["validation"] = validation
        else:
            result["fields"] = None
            result["validation"] = {"checks": [], "issues": [], "needs_review": float(doc_probs.max()) < 0.6}
        result["text"] = "\n\n".join("\n".join(b.get("text", "") for b in p["blocks"] if b.get("text")) for p in pages_out)
        result["pages"] = pages_out
        timings["assemble"] += time.time() - t
        result["timing_ms"] = {k: round(v * 1000, 1) for k, v in timings.items()}
        result["timing_ms"]["total"] = round((time.time() - t_start) * 1000, 1)
        result["model"] = {"parameters": dict(self.backend.params), "total_parameters": int(sum(self.backend.params.values()))}
        return result

	"""End-to-end document processing: file -> pages -> text segments -> labelled fields -> JSON."""

	import json
	import time
	from pathlib import Path

	import cv2
	import numpy as np

	from ..charset import NUM_CLASSES, ctc_greedy_decode
	from ..data.datasets import line_to_input, orientation_page, text_patches
	from ..features import MAX_SEGMENTS, segment_features
	from ..geometry import crop_quad, expand_quad, quad_bbox, quad_size, transform_quads
	from ..models.orientation import ANGLES, PATCH
	from ..models.recognizer import WIDTH_STRIDE
	from ..schema import DOC_TYPES, FIELD_LABELS, REL_TYPES
	from .assemble import (assemble_fields, assemble_page_structure, block_groups, block_order, groups_from_block_map,
	structure_from_groups, xy_cut)
	from .detect import DET_MAX_SIDE, detector_input, find_rules, kernels_to_quads, split_gutter_crossers
	from .inputs import estimate_skew, load_pages, rot90_matrix, rotate_image

	COMPONENTS = ("detector", "recognizer", "orientation", "extractor")

	# Frequent words used to check a proposed page rotation by reading a few lines both ways (upside-down text
	# still yields plausible-looking letters, so OCR confidence alone does not separate the two).
	COMMON_WORDS = frozenset("""
	the of and to in is for on that with as by this be are from at or an it was will not have all can has
	but which their more one were been other its into also than about after these such may when only over
	new they we you he she his her our your there two first time most some would what so if no up out who
	total tax subtotal cash change qty amount price date invoice item items paid due balance receipt number
	page table figure section chapter results data used use between each both during through under where
	""".split())
	SCHEMA_VERSION = "1.1" # 1.1: "formula" blocks; layout blocks from the model's same_block relation
	# White columns after every line the recognizer reads. Training pads each line of a batch with white to the right
	# (to a multiple of 64 px), so the model reads a line's end against white; without this strip a line read on its
	# own, or the widest line of a batch, would end at the image edge instead. 32 px covers the CNN's receptive field
	# (64 px gives identical output), and with it the output no longer depends on how lines are batched.
	RECOGNIZER_PAD = 32


	def _softmax(x, axis=-1):
	e = np.exp(x - x.max(axis, keepdims=True))
	return e / e.sum(axis, keepdims=True)


	def line_batch(lines):
	"""Recognizer input for line arrays (32, W) in [0, 1]: (batch padded with white, step mask, steps per line)."""
	wmax = max(x.shape[1] for x in lines) + RECOGNIZER_PAD
	batch = np.ones((len(lines), 1, lines[0].shape[0], wmax), np.float32)
	steps = np.array([x.shape[1] // WIDTH_STRIDE for x in lines])
	for k, x in enumerate(lines):
	batch[k, 0, :, : x.shape[1]] = x
	mask = np.arange(wmax // WIDTH_STRIDE)[None, :] >= steps[:, None]
	return batch, mask, steps


	def _sigmoid(x):
	return 1 / (1 + np.exp(-np.clip(x, -60, 60)))


	class TorchBackend:
	def __init__(self, weights_dir, device="cpu"):
	import torch

	from ..models import build, count_params

	self.torch, self.device, self.models, self.params = torch, torch.device(device), {}, {}
	for name in COMPONENTS:
	p = Path(weights_dir) / name / "best.pt"
	if not p.exists():
	raise FileNotFoundError(f"missing {p}; train it with `python -m tinydoc.train {name}`")
	ck = torch.load(p, map_location="cpu", weights_only=False)
	m = build(name, **(ck.get("arch") or {}))
	m.load_state_dict(ck["model"])
	self.models[name] = m.to(self.device).eval()
	self.params[name] = count_params(m)
	self.vocab = NUM_CLASSES

	def __call__(self, name, *inputs):
	with self.torch.no_grad():
	t = [self.torch.from_numpy(np.ascontiguousarray(x)).to(self.device) for x in inputs]
	out = self.models[name](*t)
	out = out if isinstance(out, tuple) else (out,)
	return [o.float().cpu().numpy() for o in out]


	class OnnxBackend:
	def __init__(self, onnx_dir, threads=None):
	import onnx
	import onnxruntime as ort

	opts = ort.SessionOptions()
	if threads:
	opts.intra_op_num_threads = threads
	self.sessions, self.params = {}, {}
	for name in COMPONENTS:
	p = Path(onnx_dir) / f"{name}.onnx"
	if not p.exists():
	raise FileNotFoundError(f"missing {p}; export with `python -m tinydoc.export`")
	self.sessions[name] = ort.InferenceSession(str(p), opts, providers=["CPUExecutionProvider"])
	meta = {m.key: m.value for m in onnx.load(str(p), load_external_data=False).metadata_props}
	self.params[name] = int(meta.get("trained_parameters", 0))
	# Character vocabulary the models were trained with: releases before the inline-maths characters were
	# appended have fewer classes, and their extractor must not see the newer ids.
	width = self.sessions["recognizer"].get_outputs()[0].shape[-1]
	self.vocab = width if isinstance(width, int) else NUM_CLASSES

	def __call__(self, name, *inputs):
	s = self.sessions[name]
	feed = {i.name: np.ascontiguousarray(x) for i, x in zip(s.get_inputs(), inputs)}
	return s.run(None, feed)


	class Engine:
	def __init__(self, weights="runs", backend="torch", device="cpu", recognizer_batch=None, model_blocks=None, blocks=None):
	self.backend = TorchBackend(weights, device) if backend == "torch" else OnnxBackend(weights)
	# Lines per recognizer call. On CPU (ONNX Runtime) one line per call is fastest: nothing is padded to the
	# widest line of a batch and the activations stay in cache. On 40 OmniDocBench pages (Apple M5 Pro) it cut
	# recognition from 421 to 248 ms per page at 4 threads and from 1190 to 440 ms at 1 thread, against
	# batches of 32, with the same output (see RECOGNIZER_PAD). The torch backend keeps batches, which pay
	# off on a GPU.
	self.recognizer_batch = recognizer_batch or (1 if backend == "onnx" else 32)
	# How lines are grouped into paragraphs/blocks on ordinary pages: "geometric" rules, the extractor's
	# same_block relation ("model"), or the detector's block regions ("detector"). The default comes from
	# config.json next to the weights, chosen by measuring each on validation pages.
	cfg = Path(weights) / "config.json"
	cfg = json.loads(cfg.read_text()) if cfg.exists() else {}
	if blocks is None and model_blocks is not None:
	blocks = "model" if model_blocks else "geometric"
	self.blocks = blocks or cfg.get("blocks") or ("model" if cfg.get("model_blocks") else "geometric")
	# Splitting lines merged across a column gutter: off by default (on OmniDocBench such merges were rare,
	# 5 in 380 pages, and the split hurt as often as it helped).
	self.split_gutters = cfg.get("split_gutters", False)
	# Detector input size: one pass with the long side at det_max_side. Optionally (det_fast_side) a page is
	# first detected at a smaller size and again at det_max_side only when its median line is under
	# det_min_line_px there. Off by default: on v3 it was ~15% faster and helped academic and exam pages, but
	# the official OmniDocBench evaluator showed it losing small print on newspapers (+0.017 to +0.038 text
	# edit) and magazines (+0.02), which the median-line test does not catch.
	self.det_max_side = cfg.get("det_max_side", DET_MAX_SIDE)
	self.det_fast_side = cfg.get("det_fast_side")
	self.det_min_line_px = cfg.get("det_min_line_px", 14)

	# ------------------------------------------------------------------------------------ stages
	def orientation(self, img):
	gray = orientation_page(img)
	pos = text_patches(gray, 8)
	patches = np.stack([gray[y:y + PATCH, x:x + PATCH] for x, y in pos]).astype(np.float32)[:, None] / 255
	if patches.shape[-1] != PATCH or patches.shape[-2] != PATCH:
	return 0, 1.0
	probs = _softmax(self.backend("orientation", patches)[0]).mean(0)
	k = int(probs.argmax())
	return k, float(probs[k])

	def _reading_evidence(self, img, n=16):
	"""(common-word hits, mean OCR confidence) over the n widest detected lines of an image."""
	dets = sorted(self.detect(img), key=lambda d: -quad_size(d[0])[0])[:n]
	if not dets:
	return 0, 0.0
	texts = self.recognize(img, [q for q, _ in dets])
	words = [w.strip(".,:;!?()'\"").lower() for t, _ in texts for w in t.split()]
	return sum(w in COMMON_WORDS for w in words), float(np.mean([c for _, c in texts]))

	def verify_rotation(self, img, k):
	"""Keep a proposed 90/180/270 correction only if the rotated page reads better than the original."""
	hits0, conf0 = self._reading_evidence(img)
	hitsk, confk = self._reading_evidence(np.ascontiguousarray(np.rot90(img, k)))
	if abs(hitsk - hits0) >= 2:
	return hitsk > hits0
	if abs(confk - conf0) > 0.05:
	return confk > conf0
	return True # no clear evidence either way: trust the orientation classifier

	def detect_full(self, img, blocks=True):
	"""Line detections plus the block-region map (None for detectors without a block channel, or when
	blocks=False because the caller does not use it)."""
	for side in ([self.det_fast_side] if self.det_fast_side else []) + [self.det_max_side]:
	x, scale = detector_input(img, max_side=side)
	out = self.backend("detector", x[None])[0][0]
	prob = _sigmoid(out[0])
	quads = kernels_to_quads(prob)
	if side == self.det_max_side or not quads or np.median([quad_size(q)[1] for q, _ in quads]) >= self.det_min_line_px:
	break
	block = _sigmoid(out[1]) if blocks and out.shape[0] > 1 else None
	return [(q / scale, s) for q, s in quads], block, scale

	def detect(self, img):
	return self.detect_full(img)[0]

	def recognize(self, img, quads):
	gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
	crops = []
	for q in quads:
	_, h = quad_size(q)
	crop = crop_quad(gray, expand_quad(q, 0.25 * h, 0.1 * h, 0.25 * h, 0.1 * h))
	if np.median(crop[::2, ::2]) < 100: # light text on a dark band
	crop = 255 - crop
	crops.append(line_to_input(crop))
	results = [None] * len(crops)
	order = sorted(range(len(crops)), key=lambda i: crops[i].shape[1])
	for start in range(0, len(order), self.recognizer_batch):
	idx = order[start:start + self.recognizer_batch]
	batch, mask, steps = line_batch([crops[i] for i in idx])
	probs = _softmax(self.backend("recognizer", batch, mask)[0])
	for k, i in enumerate(idx):
	text, confs = ctc_greedy_decode(probs[k, : steps[k]])
	results[i] = (text, float(np.mean(confs)) if confs else 0.0)
	return results

	def extract(self, segs, W, H):
	if not segs:
	return np.zeros((0, len(FIELD_LABELS))), np.zeros((len(REL_TYPES), 0, 0)), np.full(len(DOC_TYPES), 1 / len(DOC_TYPES))
	n = len(segs)
	labels = np.zeros((n, len(FIELD_LABELS)), np.float32)
	rel = None # sized from the model: older exports predict fewer relations
	docs = []
	# windows of MAX_SEGMENTS in reading order, overlapping by half, for very dense pages
	step = MAX_SEGMENTS // 2
	starts = list(range(0, max(n - MAX_SEGMENTS, 0) + 1, step)) or [0]
	if starts[-1] + MAX_SEGMENTS < n:
	starts.append(n - MAX_SEGMENTS)
	for s0 in starts:
	idx = list(range(s0, min(n, s0 + MAX_SEGMENTS)))
	f = segment_features([segs[i] for i in idx], W, H)
	f["char_ids"][f["char_ids"] >= getattr(self.backend, "vocab", NUM_CLASSES)] = 0
	out = self.backend("extractor", f["char_ids"][None], f["feats"][None], f["rel_boxes"][None],
	np.ones((1, len(idx)), np.float32))
	lp = _softmax(out[0][0])
	labels[idx, :lp.shape[1]] = np.maximum(labels[idx, :lp.shape[1]], lp)
	if rel is None:
	rel = np.zeros((out[1].shape[1], n, n), np.float32)
	rel[:, idx[0]:idx[-1] + 1, idx[0]:idx[-1] + 1] = np.maximum(rel[:, idx[0]:idx[-1] + 1, idx[0]:idx[-1] + 1], _sigmoid(out[1][0]))
	docs.append(_softmax(out[2][0]))
	return labels, rel, np.mean(docs, 0)

	# -------------------------------------------------------------------------------------- run
	def process(self, path, force_ocr=False):
	t_start = time.time()
	timings = {"load": 0.0, "orientation": 0.0, "deskew": 0.0, "detect": 0.0, "recognize": 0.0, "extract": 0.0, "assemble": 0.0}
	pages_out, all_segs, page_info = [], [], []
	t = time.time()
	pages = list(load_pages(path, force_ocr=force_ocr))
	timings["load"] += time.time() - t

	for page in pages:
	info = {"page": page.number, "width": page.width, "height": page.height, "text_source": page.text_source}
	if page.text_source == "pdf_text":
	segs = [{"bbox": [x0, y0, x1, y1], "quad": [[x0, y0], [x1, y0], [x1, y1], [x0, y1]], "text": t_, "conf": 1.0}
	for x0, y0, x1, y1, t_ in page.pdf_segments]
	frame_w, frame_h = page.width, page.height
	rules, block_map, det_scale = [], None, 1.0
	else:
	img = page.image
	t = time.time()
	k, kconf = self.orientation(img)
	if k != 0 and kconf > 0.6 and not self.verify_rotation(img, k):
	k, kconf = 0, 1.0 - kconf
	info["rotation_rejected"] = True
	timings["orientation"] += time.time() - t
	m_total = np.eye(3)
	if k != 0 and kconf > 0.6:
	m_total = np.vstack([rot90_matrix(k, img.shape[1], img.shape[0]), [0, 0, 1]]) @ m_total
	img = np.ascontiguousarray(np.rot90(img, k))
	info["rotation_corrected_deg"] = ANGLES[k] if (k != 0 and kconf > 0.6) else 0
	info["orientation_confidence"] = round(kconf, 3)
	t = time.time()
	skew = estimate_skew(img) # text is rotated `skew` degrees counter-clockwise
	if abs(skew) >= 0.5:
	img, m = rotate_image(img, -skew) # rotate it back
	m_total = np.vstack([m, [0, 0, 1]]) @ m_total
	info["skew_corrected_deg"] = round(skew, 2) if abs(skew) >= 0.5 else 0.0
	timings["deskew"] += time.time() - t
	t = time.time()
	dets, block_map, det_scale = self.detect_full(img, blocks=self.blocks == "detector")
	if self.split_gutters:
	dets = split_gutter_crossers(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), dets)
	timings["detect"] += time.time() - t
	t = time.time()
	texts = self.recognize(img, [q for q, _ in dets])
	timings["recognize"] += time.time() - t
	inv = np.linalg.inv(m_total)
	segs = []
	for (q, dscore), (text, conf) in zip(dets, texts):
	if not text.strip():
	continue
	orig_q = transform_quads([q], inv)[0] # report positions in the source image
	segs.append({"bbox": quad_bbox(q), "quad_src": orig_q, "text": text.strip(),
	"conf": conf * min(1.0, dscore / 0.7)})
	frame_w, frame_h = img.shape[1], img.shape[0]
	hs = [s["bbox"][3] - s["bbox"][1] for s in segs]
	rules = find_rules(img, float(np.median(hs))) if hs else []
	for i, s in enumerate(segs):
	s["page"] = page.number
	s["id"] = f"p{page.number}-s{i}"
	t = time.time()
	labels, rel, doc = self.extract(segs, frame_w, frame_h)
	timings["extract"] += time.time() - t
	boxes = [s["bbox"] for s in segs]
	groups = None
	general = bool(segs) and DOC_TYPES[int(doc.argmax())] == "general"
	if general and self.blocks == "detector" and block_map is not None:
	groups = block_order(boxes, groups_from_block_map(boxes, block_map, det_scale), rules)
	elif general and self.blocks == "model" and rel.shape[0] > REL_TYPES.index("same_block"):
	names = [FIELD_LABELS[int(r.argmax())] for r in labels]
	groups = block_order(boxes, block_groups(boxes, names, rel[REL_TYPES.index("same_block")]), rules)
	if groups is not None:
	order = [i for g in groups for i in g]
	else:
	order = xy_cut(boxes, rules=rules) if segs else []
	page_info.append((info, segs, labels, rel, doc, order, groups))

	t = time.time()
	doc_probs = np.mean([p[4] for p in page_info if len(p[1])], 0) if any(len(p[1]) for p in page_info) else np.full(len(DOC_TYPES), 1 / len(DOC_TYPES))
	doc_type = DOC_TYPES[int(doc_probs.argmax())]

	# global segment list across pages for field assembly
	offset = 0
	glob_segs, glob_labels, glob_conf, glob_order = [], [], [], []
	n_total = sum(len(p[1]) for p in page_info)
	same_row = np.zeros((n_total, n_total), np.float32)
	key_value = np.zeros((n_total, n_total), np.float32)
	for info, segs, labels, rel, _, order, groups in page_info:
	n = len(segs)
	glob_segs += segs
	glob_labels += [FIELD_LABELS[int(r.argmax())] for r in labels]
	glob_conf += [float(r.max()) for r in labels]
	glob_order += [offset + i for i in order]
	if n:
	same_row[offset:offset + n, offset:offset + n] = rel[0]
	key_value[offset:offset + n, offset:offset + n] = rel[1]
	label_names = [FIELD_LABELS[int(r.argmax())] for r in labels]
	if groups is not None:
	blocks, tables = structure_from_groups(segs, groups, label_names, rel[0])
	else:
	blocks, tables = assemble_page_structure(segs, order, label_names, rel[0] if n else np.zeros((0, 0)))
	pages_out.append({**info, "blocks": blocks, "tables": tables, "segments": [
	{"id": s["id"], "text": s["text"], "bbox": [round(v, 1) for v in (quad_bbox(s["quad_src"]) if "quad_src" in s else s["bbox"])],
	"quad": np.round(np.asarray(s.get("quad_src", s.get("quad"))), 1).tolist(),
	"ocr_confidence": round(s["conf"], 3), "label": label_names[i], "label_confidence": round(float(labels[i].max()), 3)}
	for i, s in enumerate(segs)]})
	offset += n

	result = {
	"schema_version": SCHEMA_VERSION,
	"source": {"file": str(Path(path).name), "pages": len(pages)},
	"document_type": {"value": doc_type, "confidence": round(float(doc_probs.max()), 3),
	"scores": {d: round(float(p), 3) for d, p in zip(DOC_TYPES, doc_probs)}},
	}
	if n_total == 0:
	result["document_type"] = {"value": None, "confidence": 0.0, "scores": None}
	result["fields"] = None
	result["validation"] = {"checks": [], "issues": ["no text found"], "needs_review": True}
	elif doc_type in ("receipt", "invoice"):
	fields, validation = assemble_fields(glob_segs, glob_order, glob_labels, glob_conf, same_row, key_value)
	result["fields"] = fields
	result["validation"] = validation
	else:
	result["fields"] = None
	result["validation"] = {"checks": [], "issues": [], "needs_review": float(doc_probs.max()) < 0.6}
	result["text"] = "\n\n".join("\n".join(b.get("text", "") for b in p["blocks"] if b.get("text")) for p in pages_out)
	result["pages"] = pages_out
	timings["assemble"] += time.time() - t
	result["timing_ms"] = {k: round(v * 1000, 1) for k, v in timings.items()}
	result["timing_ms"]["total"] = round((time.time() - t_start) * 1000, 1)
	result["model"] = {"parameters": dict(self.backend.params), "total_parameters": int(sum(self.backend.params.values()))}
	return result