File size: 2,366 Bytes
1b4ed07 7fb4118 1b4ed07 | 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 | # 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())
|