File size: 2,741 Bytes
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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | # Copyright 2026 AutomatosX
# SPDX-License-Identifier: AGPL-3.0-only
# Modified 2026-10-17 for v1.2; prior component notices are in NOTICE.
"""``ax-doclayout`` command line: one page raster in, canonical JSON out."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from ax_doclayout.contract import PageImage
from ax_doclayout.errors import AXDocLayoutError, InvalidPageImageError
from ax_doclayout.model import BACKENDS, AXDocLayout
_PDF_MAGIC = b"%PDF-"
def _load_rgb(path: Path):
"""Decode a page raster; PDFs are rejected by content and by file name."""
try:
with path.open("rb") as stream:
head = stream.read(len(_PDF_MAGIC))
except OSError as error:
raise InvalidPageImageError(
f"cannot read page raster: {error.strerror}"
) from None
if head.startswith(_PDF_MAGIC) or path.suffix.lower() == ".pdf":
raise InvalidPageImageError("the CLI takes a decoded page raster, not a PDF")
import numpy as np
from PIL import Image, UnidentifiedImageError
try:
with Image.open(path) as image:
return np.asarray(image.convert("RGB"), dtype=np.uint8)
except (OSError, UnidentifiedImageError, Image.DecompressionBombError, ValueError) as error:
raise InvalidPageImageError(
f"cannot decode page raster: {type(error).__name__}"
) from None
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(prog="ax-doclayout")
commands = parser.add_subparsers(dest="command", required=True)
predict = commands.add_parser(
"predict", help="detect layout regions on one page raster"
)
predict.add_argument("image", type=Path)
predict.add_argument("--page-index", type=int, required=True)
predict.add_argument("--model", required=True)
predict.add_argument("--revision", required=True)
predict.add_argument("--backend", choices=BACKENDS, required=True)
predict.add_argument("--output", default="-")
arguments = parser.parse_args(argv)
try:
page = PageImage(
rgb=_load_rgb(arguments.image), page_index=arguments.page_index
)
model = AXDocLayout.from_pretrained(
arguments.model, revision=arguments.revision, backend=arguments.backend
)
payload = json.dumps(model.predict(page).to_json(), indent=2)
except AXDocLayoutError as error:
print(f"ax-doclayout: {error}", file=sys.stderr)
return 2
if arguments.output == "-":
print(payload)
else:
Path(arguments.output).write_text(payload + "\n", encoding="utf-8")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|