| |
| |
| |
| """``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()) |
|
|
