teklia-layout-reader 0.2.1__py3-none-any.whl

layout_reader/inference.py

@@ -0,0 +1,215 @@
+import json
+import logging
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+from layout_reader.helpers import (
+    boxes_to_inputs,
+    load_dataset_split,
+    load_model,
+    parse_logits,
+    prepare_inputs,
+    save_visualization,
+    sort_sample,
+)
+
+logger = logging.getLogger(__name__)
+DATA_FILES = ["train", "val", "dev", "test"]
+
+
+@dataclass
+class PageResult:
+    boxes: list[Any]  # List of bounding boxes to reorder
+    classes: list[Any]  # List of corresponding classes
+    separators: list[Any]  # List of separators, used as context by the model
+    predicted_order: list[int]  # Detected reading order
+
+    # Optional evaluation outputs
+    target_order: list[int] | None = None  # Target reading order
+
+    # Optional visualization output
+    visualization: str | None = None  # Path to the visualization image
+
+    @property
+    def average_relative_distance(self) -> float:
+        """
+        Compute the Average Relative Distance (ARD) between predicted and ground-truth ordering.
+        """
+        gt_zones = [self.boxes[i] for i in self.target_order]
+        predicted_zones = [self.boxes[i] for i in self.predicted_order]
+
+        ard = 0
+
+        if not gt_zones and not predicted_zones:
+            return 0
+
+        if not gt_zones or not predicted_zones:
+            return float(max(len(gt_zones), len(predicted_zones)))
+
+        for pred_idx, pred_zone in enumerate(predicted_zones):
+            if pred_zone in gt_zones:
+                ard += abs(gt_zones.index(pred_zone) - pred_idx)
+            else:
+                ard += len(predicted_zones)
+
+        return ard / len(predicted_zones)
+
+    def to_json(self) -> dict[str, Any]:
+        data = asdict(self)
+
+        if self.target_order is not None:
+            data["average_relative_distance"] = self.average_relative_distance
+
+        return data
+
+
+def add_predict_command(subcommands) -> None:
+    parser = subcommands.add_parser(
+        "inference",
+        description="Predict and evaluate a model on a Layout Reader dataset.",
+        help="Predict and evaluate a model on a Layout Reader dataset.",
+    )
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        help="Path to the local LayoutReader dataset directory. The directory must contain .jsonl.gz files.",
+    )
+    parser.add_argument(
+        "--split",
+        choices=DATA_FILES,
+        help="Dataset split to use. Must match the name of a corresponding archive in the --dataset directory (e.g., 'train' -> train.jsonl.gz).",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        help="Name of the LayoutReader checkpoint dataset",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        help="Output directory where results will be saved. If --visualization, it will also plot the predicted reading order on each image.",
+        required=True,
+    )
+    parser.add_argument(
+        "--with-classes",
+        action="store_true",
+        help="Whether to use the zone classes for the prediction.",
+    )
+    parser.add_argument(
+        "--with-separators",
+        action="store_true",
+        help="Whether to use the separators for the prediction.",
+    )
+    parser.add_argument(
+        "--sort-method",
+        choices=["sortxy_by_column", "sortxy", "sortyx", "random"],
+        default=None,
+        help="How to pre-order the input zones. By default, no sorting method will be applied.",
+    )
+    parser.add_argument(
+        "--images",
+        type=Path,
+        help="Path to the images (required to visualize predictions).",
+        default=None,
+    )
+    parser.add_argument(
+        "--visualize",
+        action="store_true",
+        help="Whether to visualize the predicted reading order. Plots will be saved in output_dir.",
+    )
+    parser.set_defaults(func=run)
+
+
+def predict(model, boxes, classes, separators) -> list[int]:
+    if not boxes:
+        return []
+    inputs = boxes_to_inputs(boxes, cls=classes, separators=separators)
+    inputs = prepare_inputs(inputs, model)
+    logits = model(**inputs).logits.squeeze(0)
+    predicted_orders = parse_logits(logits, len(boxes))
+    return predicted_orders
+
+
+def process_page(
+    page,
+    model,
+    with_classes: bool,
+    with_separators: bool,
+    visualize: bool,
+    images: Path | None,
+    output_dir: Path,
+    sort_method: str | None,
+) -> PageResult:
+    # Pre-order boxes and classes
+    if sort_method is not None:
+        page = sort_sample(page, sort_ratio=1.0, sort_method=sort_method)
+
+    boxes = page["source_boxes"]
+    classes = page["source_classes"] if with_classes else []
+    separators = page["separators"] if with_separators else []
+    predicted_order = predict(model, boxes, classes, separators)
+
+    result = PageResult(
+        boxes=boxes,
+        classes=classes,
+        separators=separators,
+        predicted_order=predicted_order,
+    )
+
+    # Run evaluation if the ground truth is available
+    if "target_index" in page:
+        result.target_order = [i - 1 for i in page["target_index"]]
+
+    if visualize and images is not None and result.boxes:
+        image_path = (images / page["sample_id"]).with_suffix(".jpg")
+        output_path = output_dir / f"{page['sample_id']}_visualize.jpg"
+        save_visualization(
+            image_path=image_path,
+            boxes=boxes,
+            predicted_order=predicted_order,
+            output_path=output_path,
+        )
+
+        result.visualization = str(output_path)
+    return result
+
+
+def run(
+    dataset: str,
+    split: str,
+    model: str,
+    output_dir: Path,
+    with_classes: bool = False,
+    with_separators: bool = False,
+    images: Path | None = None,
+    visualize: bool = False,
+    sort_method: str | None = None,
+) -> None:
+    logger.info(f"Loading model {model}")
+    model = load_model(model_path=model)
+
+    logger.info(f"Loading dataset {dataset}/{split}.jsonl.gz")
+    dataset = load_dataset_split(dataset, split)
+
+    results = {}
+
+    if visualize and images is None:
+        logger.warning("Skipping visualization, as --images is not defined")
+
+    for page in dataset:
+        page_result = process_page(
+            page=page,
+            model=model,
+            with_classes=with_classes,
+            with_separators=with_separators,
+            visualize=visualize,
+            images=images,
+            output_dir=output_dir,
+            sort_method=sort_method,
+        )
+
+        results[page["sample_id"]] = page_result.to_json()
+
+    output_dir.mkdir(exist_ok=True)
+    (output_dir / "predictions.json").write_text(json.dumps(results, indent=4))

layout_reader/train/sft.py

@@ -0,0 +1,69 @@
+import logging
+from pathlib import Path
+
+from transformers import (
+    LayoutLMv3ForTokenClassification,
+)
+from transformers.trainer import Trainer
+from trl import SFTConfig
+
+from layout_reader.helpers import (
+    MAX_LEN,
+    DataCollator,
+    load_dataset_split,
+    read_yaml,
+    sort_sample,
+)
+
+logger = logging.getLogger()
+
+
+def run_training(config: dict):
+    train_dataset = load_dataset_split(Path(config["dataset_dir"]), "train")
+    dev_dataset = load_dataset_split(Path(config["dataset_dir"]), "dev")
+    if config["sort_ratio"] > 0:
+
+        def apply_sort_fn(dataset):
+            return sort_sample(
+                dataset,
+                sort_ratio=config["sort_ratio"],
+                sort_method=config["sort_method"],
+            )
+
+        train_dataset = train_dataset.map(apply_sort_fn)
+        dev_dataset = dev_dataset.map(apply_sort_fn)
+
+    logger.info(
+        f"Train dataset size: {len(train_dataset)}, Dev dataset size: {len(dev_dataset)}"
+    )
+
+    model = LayoutLMv3ForTokenClassification.from_pretrained(
+        config["model_dir"], num_labels=MAX_LEN, visual_embed=False
+    )
+
+    data_collator = DataCollator(
+        with_classes=config["with_classes"], with_separators=config["with_separators"]
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=SFTConfig(**config["SFTTrainer"]),
+        train_dataset=train_dataset,
+        eval_dataset=dev_dataset,
+        data_collator=data_collator,
+    )
+    trainer.train()
+
+
+def add_training_command(subcommands) -> None:
+    parser = subcommands.add_parser(
+        "train",
+        description="Train a Layout Reader model on a dataset.",
+        help="Train a Layout Reader model on a dataset.",
+    )
+    parser.add_argument(
+        "--config",
+        type=read_yaml,
+        help="Name of the LayoutReader dataset",
+    )
+    parser.set_defaults(func=run_training)

teklia_layout_reader-0.2.1.dist-info/METADATA

@@ -0,0 +1,62 @@
+Metadata-Version: 2.4
+Name: teklia-layout-reader
+Version: 0.2.1
+Summary: Scripts for Layout Reader
+Author-email: Teklia <contact@teklia.com>
+Maintainer-email: Teklia <contact@teklia.com>
+License-Expression: MIT
+Keywords: python
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: teklia-toolbox==0.1.11
+Requires-Dist: datasets==3.2.0
+Requires-Dist: mdutils==1.6.0
+Requires-Dist: prettytable==3.13.0
+Requires-Dist: opencv-python-headless==4.12.0.88
+Requires-Dist: matplotlib==3.10.6
+Requires-Dist: deepspeed==0.18.1
+Requires-Dist: trl==0.23.0
+Requires-Dist: mpi4py==4.1.1
+Requires-Dist: wandb==0.23.1
+Requires-Dist: transformers==4.57.6
+Requires-Dist: colour==0.1.5
+Requires-Dist: accelerate==1.12.0
+
+# Layout Reader
+
+Scripts for Layout Reader
+
+### Development
+
+For development and tests purpose it may be useful to install the project as a editable package with pip.
+
+* Use a virtualenv (e.g. with virtualenvwrapper `mkvirtualenv -a . layout-reader`)
+* Install layout-reader as a package (e.g. `pip install -e .`)
+
+### Linter
+
+Code syntax is analyzed before submitting the code.\
+To run the linter tools suite you may use pre-commit.
+
+```shell
+pip install pre-commit
+pre-commit run -a
+```
+
+### Run tests
+
+Tests are executed with `tox` using [pytest](https://pytest.org).
+
+```shell
+pip install tox
+tox
+```
+
+To recreate tox virtual environment (e.g. a dependencies update), you may run `tox -r`.
+
+Run a single test module: `tox -- <test_path>`
+Run a single test: `tox -- <test_path>::<test_function>`

teklia_layout_reader-0.2.1.dist-info/RECORD

@@ -0,0 +1,22 @@
+layout_reader/__init__.py,sha256=5doGcBcb3kTQMoDbjiESh1g0CrG_mHTJlGRX4Uo7NRU,156
+layout_reader/cli.py,sha256=4M4nzV7kZmF-LmUqHie_eeXJIhvHdhgAi-_q9gm7yEo,665
+layout_reader/helpers.py,sha256=rWJT0glBqOYCj4tvktiPunxoTHiNb5v2oD5MainDrhU,11675
+layout_reader/inference.py,sha256=ECujeVPfYI6DM0yl2WorzmWwD5bMZTblm4uFg2ggldY,6616
+layout_reader/datasets/__init__.py,sha256=ppjB6O0uDJLR3FG3aKnZmtZzYLWJC8Aeb3KjnAWawzE,2663
+layout_reader/datasets/analyze.py,sha256=Lw_Coqpudlz2v4quZPvWr61AEJYTkSYJNUHkzeCxZ38,4700
+layout_reader/datasets/extract.py,sha256=Sdoe8kMBrjKd2ydGYDySc-6e2w6_0BgFa8sLqTsQNKU,8726
+layout_reader/datasets/lsd.py,sha256=jVE6ByzR9OPnvLvpYbeF5adDWHEeVLFor8-9o9uVzGE,4211
+layout_reader/datasets/utils.py,sha256=Nvw3y441WS-fU-9qXCMcM3_UZaEBnBsS19Is0zu6Z_w,3393
+layout_reader/train/sft.py,sha256=yb2MqEZhQDxkBcUXqvPC-ebjG7BuaLzoT2vEaRo7HlE,1872
+tests/__init__.py,sha256=Z6xMdy1fM-m3Ncq4LMozIv_LvHsHnUCcFPGe7qaI3h0,78
+tests/conftest.py,sha256=CX31U3NxXMLQp28DK2uU4GELObTAKdaZvlX9B6JlGrY,500
+tests/test_analyze.py,sha256=DeOQ2vy0qYwEmIquhLCN_LrOgVwv_sO55apXaLpYGxo,405
+tests/test_cli.py,sha256=vizFW9k23hKZQSNjXZOXGtwfzS646HK8k4Huag4T1lI,246
+tests/test_extract.py,sha256=ruT3fedzXCQ1FdW3uYYEZA2VdX3GxK67hzdoLi1u9jA,3251
+tests/test_helpers.py,sha256=sOX3JWfZE_EL2ZD6qNF9vFbqSm6Ak2158nZcFsdD91o,13896
+tests/test_predict.py,sha256=DoZZtG8AlJtwxF2f6bz1CPCIrMjLjs4MbHRA8XF4Q5k,1613
+teklia_layout_reader-0.2.1.dist-info/METADATA,sha256=DU5rASxjSgEjEE0LZBTIOvSxgLjUBZM0owXPyAQ5kr8,1735
+teklia_layout_reader-0.2.1.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+teklia_layout_reader-0.2.1.dist-info/entry_points.txt,sha256=9kJyzLJuIhGeoYnwHxYITSNh3--7v7HTMfgwIpG0_fc,64
+teklia_layout_reader-0.2.1.dist-info/top_level.txt,sha256=G0dTKLPB9MCL_4pPL8Y0aqjT6ujCWb3Y1Si3UutIvCU,20
+teklia_layout_reader-0.2.1.dist-info/RECORD,,

teklia_layout_reader-0.2.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.10.2)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

teklia_layout_reader-0.2.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+teklia-layout-reader = layout_reader.cli:main

teklia_layout_reader-0.2.1.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+layout_reader
+tests

tests/__init__.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+FIXTURES = Path(__file__).resolve().parent / "data"

tests/conftest.py

@@ -0,0 +1,19 @@
+import os
+
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def _setup_environment(responses):
+    """Setup needed environment variables"""
+
+    # Allow accessing remote API schemas
+    # defaulting to the prod environment
+    schema_url = os.environ.get(
+        "ARKINDEX_API_SCHEMA_URL",
+        "https://arkindex.teklia.com/api/v1/openapi/?format=openapi-json",
+    )
+    responses.add_passthru(schema_url)
+
+    # Set schema url in environment
+    os.environ["ARKINDEX_API_SCHEMA_URL"] = schema_url

tests/test_analyze.py

@@ -0,0 +1,14 @@
+from layout_reader.datasets.analyze import Statistics
+from tests import FIXTURES
+
+DATASET = FIXTURES / "lr_dataset" / "train.jsonl.gz"
+REPORT = FIXTURES / "dataset_report.md"
+
+
+def test_run(tmp_path) -> None:
+    output_file = tmp_path / "report.md"
+    stats = Statistics(output=str(output_file))
+    stats.run(
+        filenames=[DATASET],
+    )
+    assert output_file.read_text() == REPORT.read_text()

tests/test_cli.py

@@ -0,0 +1,11 @@
+import importlib
+
+
+def test_dummy():
+    assert True
+
+
+def test_import():
+    """Import our newly created module, through importlib to avoid parsing issues"""
+    cli = importlib.import_module("layout_reader.cli")
+    assert hasattr(cli, "main")

tests/test_extract.py

@@ -0,0 +1,130 @@
+import random
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from layout_reader.datasets.extract import (
+    Mode,
+    run_dataset_extraction,
+)
+from layout_reader.datasets.utils import load_gzip_jsonl
+from tests import FIXTURES
+
+
+@pytest.fixture
+def SPLIT() -> str:
+    return "train"
+
+
+@pytest.fixture
+def hf_dataset() -> str:
+    return str(FIXTURES / "hf_dataset")
+
+
+@pytest.fixture
+def yolo_dataset() -> str:
+    return str(FIXTURES / "yolo_dataset")
+
+
+@pytest.fixture
+def lr_dataset() -> str:
+    return str(FIXTURES / "lr_dataset")
+
+
+@pytest.mark.parametrize(
+    ("shuffle_rate", "extract_classes", "extract_separators"),
+    [
+        (0.0, False, False),
+        (0.0, True, False),
+        (0.0, False, True),
+        (1.0, True, True),
+        (0.5, False, True),
+        (0.5, True, False),
+    ],
+)
+def test_run_dataset_extraction_modes(
+    hf_dataset,
+    yolo_dataset,
+    shuffle_rate,
+    extract_classes,
+    extract_separators,
+    tmp_path,
+):
+    output_dir_hf = tmp_path / "output_hf"
+    output_dir_yolo = tmp_path / "output_yolo"
+
+    random.seed(42)
+    np.random.seed(42)
+    run_dataset_extraction(
+        hf_dataset,
+        Mode.HF,
+        shuffle_rate=shuffle_rate,
+        output_dir=output_dir_hf,
+        extract_classes=extract_classes,
+        extract_separators=extract_separators,
+    )
+
+    random.seed(42)
+    np.random.seed(42)
+    run_dataset_extraction(
+        yolo_dataset,
+        Mode.YOLO,
+        shuffle_rate=shuffle_rate,
+        output_dir=output_dir_yolo,
+        extract_classes=extract_classes,
+        extract_separators=extract_separators,
+    )
+    assert (Path(output_dir_hf) / "train.jsonl.gz").exists()
+    assert (Path(output_dir_yolo) / "train.jsonl.gz").exists()
+    assert load_gzip_jsonl(Path(output_dir_hf) / "train.jsonl.gz") == load_gzip_jsonl(
+        Path(output_dir_yolo) / "train.jsonl.gz"
+    )
+
+
+@pytest.mark.parametrize(
+    ("shuffle_rate", "extract_classes", "extract_separators"),
+    [
+        (0.5, True, True),
+    ],
+)
+def test_run_dataset_extraction(
+    hf_dataset,
+    yolo_dataset,
+    lr_dataset,
+    shuffle_rate,
+    extract_classes,
+    extract_separators,
+    tmp_path,
+):
+    output_dir_hf = tmp_path / "output_hf"
+    output_dir_yolo = tmp_path / "output_yolo"
+
+    random.seed(42)
+    np.random.seed(42)
+    run_dataset_extraction(
+        hf_dataset,
+        Mode.HF,
+        shuffle_rate=shuffle_rate,
+        output_dir=output_dir_hf,
+        extract_classes=extract_classes,
+        extract_separators=extract_separators,
+    )
+
+    random.seed(42)
+    np.random.seed(42)
+    run_dataset_extraction(
+        yolo_dataset,
+        Mode.YOLO,
+        shuffle_rate=shuffle_rate,
+        output_dir=output_dir_yolo,
+        extract_classes=extract_classes,
+        extract_separators=extract_separators,
+    )
+    assert (Path(output_dir_hf) / "train.jsonl.gz").exists()
+    assert (Path(output_dir_yolo) / "train.jsonl.gz").exists()
+
+    dataset_from_hf = load_gzip_jsonl(Path(output_dir_hf) / "train.jsonl.gz")
+    dataset_from_yolo = load_gzip_jsonl(Path(output_dir_yolo) / "train.jsonl.gz")
+    expected_dataset = load_gzip_jsonl(Path(lr_dataset) / "train.jsonl.gz")
+    assert dataset_from_hf == dataset_from_yolo == expected_dataset