inference/examples/predict_page.py

# Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""Detect layout on one page raster using the pinned public Hub weights."""

from __future__ import annotations

import argparse
import json
import sys
import warnings
from pathlib import Path

import numpy as np
from PIL import Image

from ax_doclayout import AXDocLayout, AXDocLayoutError, InvalidPageImageError, PageImage
from ax_doclayout.contract import MAX_PAGE_PIXELS

MODEL_ID = "AutomatosX/AXDocLayout-M-v1.2-CUDA-TF32"
REVISION = "7d2e54941a3c2c9d1ac4af27f21531bf39dfb1dd"


def load_page(path: Path, *, page_index: int) -> PageImage:
    """Decode a raster without recording its filename or pixels in the output."""

    try:
        with path.open("rb") as handle:
            if handle.read(5) == b"%PDF-" or path.suffix.lower() == ".pdf":
                raise InvalidPageImageError("provide a page raster image")
        with warnings.catch_warnings():
            warnings.simplefilter("error", Image.DecompressionBombWarning)
            with Image.open(path) as image:
                if image.width * image.height > MAX_PAGE_PIXELS:
                    raise InvalidPageImageError("page raster exceeds pixel limit")
                rgb = np.asarray(image.convert("RGB"), dtype=np.uint8)
    except (OSError, ValueError, Image.DecompressionBombError, Image.DecompressionBombWarning):
        raise InvalidPageImageError("page raster could not be decoded") from None
    return PageImage(rgb=rgb, page_index=page_index)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("image", type=Path)
    parser.add_argument("--backend", choices=("onnx", "pytorch"), default="onnx")
    parser.add_argument("--page-index", type=int, default=0)
    args = parser.parse_args(argv)
    try:
        page = load_page(args.image, page_index=args.page_index)
        model = AXDocLayout.from_pretrained(
            MODEL_ID, revision=REVISION, backend=args.backend,
        )
        print(json.dumps(model.predict(page).to_json(), indent=2, allow_nan=False))
    except AXDocLayoutError as error:
        print(f"ax-doclayout: {error}", file=sys.stderr)
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())