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

tests/test_helpers.py

@@ -0,0 +1,438 @@
+import pytest
+import torch
+
+from layout_reader.helpers import (
+    CLS_TOKEN_ID,
+    EOS_TOKEN_ID,
+    IGNORE_LABEL_ID,
+    PAD_TOKEN_ID,
+    DataCollator,
+    boxes_to_inputs,
+    load_dataset_split,
+    parse_logits,
+    sort_sample,
+)
+from tests import FIXTURES
+
+
+@pytest.fixture
+def dataset() -> str:
+    return FIXTURES / "lr_dataset"
+
+
+@pytest.mark.parametrize(
+    (
+        "features",
+        "with_classes",
+        "with_separators",
+        "expected_boxes",
+        "expected_input_ids",
+        "expected_attention_masks",
+        "expected_labels",
+    ),
+    [
+        (
+            [
+                {
+                    "source_boxes": [[1, 1, 2, 2], [2, 2, 3, 3]],
+                    "separators": [[1, 1, 1, 6]],
+                    "target_index": [1, 2],
+                    "source_classes": [10, 11],
+                }
+            ],
+            True,
+            True,
+            torch.tensor(
+                [[[0, 0, 0, 0], [1, 1, 2, 2], [2, 2, 3, 3], [1, 1, 1, 6], [0, 0, 0, 0]]]
+            ),
+            torch.tensor([[CLS_TOKEN_ID, 10, 11, PAD_TOKEN_ID, EOS_TOKEN_ID]]),
+            torch.tensor([[1, 1, 1, 1, 1]]),
+            torch.tensor([[IGNORE_LABEL_ID, 0, 1, IGNORE_LABEL_ID, IGNORE_LABEL_ID]]),
+        ),
+        (
+            [
+                {
+                    "source_boxes": [[1, 1, 2, 2], [2, 2, 3, 3]],
+                    "separators": [[1, 1, 1, 6]],
+                    "target_index": [1, 2],
+                    "source_classes": [10, 11],
+                }
+            ],
+            False,
+            True,
+            torch.tensor(
+                [[[0, 0, 0, 0], [1, 1, 2, 2], [2, 2, 3, 3], [1, 1, 1, 6], [0, 0, 0, 0]]]
+            ),
+            torch.tensor(
+                [[CLS_TOKEN_ID, PAD_TOKEN_ID, PAD_TOKEN_ID, PAD_TOKEN_ID, EOS_TOKEN_ID]]
+            ),
+            torch.tensor([[1, 1, 1, 1, 1]]),
+            torch.tensor([[IGNORE_LABEL_ID, 0, 1, IGNORE_LABEL_ID, IGNORE_LABEL_ID]]),
+        ),
+        (
+            [
+                {
+                    "source_boxes": [[1, 1, 2, 2], [2, 2, 3, 3]],
+                    "separators": [[1, 1, 1, 6]],
+                    "target_index": [1, 2],
+                    "source_classes": [10, 11],
+                }
+            ],
+            False,
+            False,
+            torch.tensor([[[0, 0, 0, 0], [1, 1, 2, 2], [2, 2, 3, 3], [0, 0, 0, 0]]]),
+            torch.tensor([[CLS_TOKEN_ID, PAD_TOKEN_ID, PAD_TOKEN_ID, EOS_TOKEN_ID]]),
+            torch.tensor([[1, 1, 1, 1]]),
+            torch.tensor([[IGNORE_LABEL_ID, 0, 1, IGNORE_LABEL_ID]]),
+        ),
+        (
+            [
+                {
+                    "source_boxes": [[1, 1, 2, 2], [2, 2, 3, 3]],
+                    "separators": [[1, 1, 1, 6]],
+                    "target_index": [1, 2],
+                    "source_classes": [10, 11],
+                }
+            ],
+            True,
+            False,
+            torch.tensor([[[0, 0, 0, 0], [1, 1, 2, 2], [2, 2, 3, 3], [0, 0, 0, 0]]]),
+            torch.tensor([[CLS_TOKEN_ID, 10, 11, EOS_TOKEN_ID]]),
+            torch.tensor([[1, 1, 1, 1]]),
+            torch.tensor([[IGNORE_LABEL_ID, 0, 1, IGNORE_LABEL_ID]]),
+        ),
+    ],
+)
+def test_data_collator(
+    features,
+    with_classes,
+    with_separators,
+    expected_labels,
+    expected_boxes,
+    expected_input_ids,
+    expected_attention_masks,
+):
+    collator = DataCollator(with_classes=with_classes, with_separators=with_separators)
+    output = collator(features)
+
+    assert set(output.keys()) == {
+        "bbox",
+        "attention_mask",
+        "labels",
+        "input_ids",
+    }
+
+    assert torch.equal(output["labels"], expected_labels)
+    assert torch.equal(output["bbox"], expected_boxes)
+    assert torch.equal(output["input_ids"], expected_input_ids)
+    assert torch.equal(output["attention_mask"], expected_attention_masks)
+
+
+@pytest.mark.parametrize(
+    ("boxes", "classes", "separators", "expected_output"),
+    [
+        (
+            [[1, 1, 2, 2], [2, 2, 3, 3]],
+            [10, 11],
+            [[1, 1, 1, 6]],
+            {
+                "bbox": torch.tensor(
+                    [
+                        [
+                            [0, 0, 0, 0],
+                            [1, 1, 2, 2],
+                            [2, 2, 3, 3],
+                            [1, 1, 1, 6],
+                            [0, 0, 0, 0],
+                        ]
+                    ]
+                ),
+                "attention_mask": torch.tensor([[1, 1, 1, 1, 1]]),
+                "input_ids": torch.tensor(
+                    [[CLS_TOKEN_ID, 10, 11, PAD_TOKEN_ID, EOS_TOKEN_ID]]
+                ),
+            },
+        ),
+        (
+            [[1, 1, 2, 2], [2, 2, 3, 3]],
+            [],
+            [[1, 1, 1, 6]],
+            {
+                "bbox": torch.tensor(
+                    [
+                        [
+                            [0, 0, 0, 0],
+                            [1, 1, 2, 2],
+                            [2, 2, 3, 3],
+                            [1, 1, 1, 6],
+                            [0, 0, 0, 0],
+                        ]
+                    ]
+                ),
+                "attention_mask": torch.tensor([[1, 1, 1, 1, 1]]),
+                "input_ids": torch.tensor(
+                    [
+                        [
+                            CLS_TOKEN_ID,
+                            PAD_TOKEN_ID,
+                            PAD_TOKEN_ID,
+                            PAD_TOKEN_ID,
+                            EOS_TOKEN_ID,
+                        ]
+                    ]
+                ),
+            },
+        ),
+        (
+            [[1, 1, 2, 2], [2, 2, 3, 3]],
+            [],
+            [],
+            {
+                "bbox": torch.tensor(
+                    [[[0, 0, 0, 0], [1, 1, 2, 2], [2, 2, 3, 3], [0, 0, 0, 0]]]
+                ),
+                "attention_mask": torch.tensor([[1, 1, 1, 1]]),
+                "input_ids": torch.tensor(
+                    [[CLS_TOKEN_ID, PAD_TOKEN_ID, PAD_TOKEN_ID, EOS_TOKEN_ID]]
+                ),
+            },
+        ),
+    ],
+)
+def test_boxes_to_input(boxes, classes, separators, expected_output):
+    output = boxes_to_inputs(boxes, classes, separators)
+
+    assert set(output.keys()) == {
+        "bbox",
+        "attention_mask",
+        "input_ids",
+    }
+    assert torch.equal(output["attention_mask"], expected_output["attention_mask"])
+    assert torch.equal(output["bbox"], expected_output["bbox"])
+    assert torch.equal(output["input_ids"], expected_output["input_ids"])
+
+
+@pytest.mark.parametrize(
+    ("logits", "length", "expected_order"),
+    [
+        (
+            torch.tensor(
+                [
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # CLS token
+                    [0.1, 0.8, 0.05, 0.05, 0.0],  # Element 0 - position 0
+                    [0.7, 0.1, 0.05, 0.1, 0.05],  # Element 1 - ignored
+                    [0.05, 0.25, 0.1, 0.55, 0.05],  # Element 2 - ignored
+                    [0.05, 0.15, 0.7, 0.05, 0.05],  # Element 3 - ignored
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # EOS token
+                ]
+            ),
+            1,
+            [0],
+        ),
+        (
+            torch.tensor(
+                [
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # CLS token
+                    [0.1, 0.8, 0.05, 0.05, 0.0],  # Element 0 - position 1
+                    [0.7, 0.1, 0.05, 0.1, 0.05],  # Element 1 - position 0
+                    [0.05, 0.25, 0.1, 0.55, 0.05],  # Element 2 - position 2
+                    [0.05, 0.15, 0.7, 0.05, 0.05],  # Element 3 - ignored
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # EOS token
+                ]
+            ),
+            3,
+            [1, 0, 2],
+        ),
+        (
+            torch.tensor(
+                [
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # CLS token
+                    [0.1, 0.8, 0.05, 0.05, 0.0],  # Element 0 - position 1
+                    [0.7, 0.1, 0.05, 0.1, 0.05],  # Element 1 - position 0
+                    [0.05, 0.25, 0.1, 0.55, 0.05],  # Element 2 - position 3
+                    [0.05, 0.15, 0.7, 0.05, 0.05],  # Element 3 - position 2
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # EOS token
+                ]
+            ),
+            4,
+            [1, 0, 3, 2],
+        ),
+        (
+            torch.tensor(
+                [  # Conflict element 0 and 1 (not equal)
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # CLS token
+                    [
+                        0.6,
+                        0.2,
+                        0.05,
+                        0.1,
+                        0.05,
+                    ],  # Element 0 - position 0 preferred (conflict -> position 1)
+                    [
+                        0.7,
+                        0.1,
+                        0.05,
+                        0.1,
+                        0.05,
+                    ],  # Element 1 - position 0 preferred (conflict -> position 0)
+                    [0.05, 0.25, 0.1, 0.55, 0.05],  # Element 2 - position 3 preferred
+                    [0.05, 0.15, 0.7, 0.05, 0.05],  # Element 3 - position 2 preferred
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # EOS token
+                ]
+            ),
+            4,
+            [1, 0, 3, 2],
+        ),
+        (
+            torch.tensor(
+                [  # Conflict element 0 and 1 (equal)
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # CLS token
+                    [
+                        0.7,
+                        0.1,
+                        0.05,
+                        0.1,
+                        0.05,
+                    ],  # Element 0 - position 0 preferred (conflict -> position 1)
+                    [
+                        0.7,
+                        0.1,
+                        0.05,
+                        0.1,
+                        0.05,
+                    ],  # Element 1 - position 0 preferred (conflict -> position 0)
+                    [0.05, 0.25, 0.1, 0.55, 0.05],  # Element 2 - position 3 preferred
+                    [0.05, 0.15, 0.7, 0.05, 0.05],  # Element 3 - position 2 preferred
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # EOS token
+                ]
+            ),
+            4,
+            [0, 1, 3, 2],
+        ),
+        (
+            torch.tensor(
+                [  # Cascade conflicts
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # CLS token
+                    [
+                        0.6,
+                        0.1,
+                        0.05,
+                        0.1,
+                        0.05,
+                    ],  # Element 0 - position 0 preferred (conflict elements #1 #2 #3 -> position 3)
+                    [
+                        0.7,
+                        0.1,
+                        0.05,
+                        0.1,
+                        0.05,
+                    ],  # Element 1 - position 0 preferred (conflict element #0 -> position 0)
+                    [
+                        0.05,
+                        0.25,
+                        0.55,
+                        0.1,
+                        0.05,
+                    ],  # Element 2 - position 2 preferred (conflict element #0 -> position 2)
+                    [
+                        0.0,
+                        0.7,
+                        0.3,
+                        0.05,
+                        0.05,
+                    ],  # Element 3 - position 1 preferred (conflict element #0 -> position 1)
+                    [0.2, 0.2, 0.2, 0.2, 0.2],  # EOS token
+                ]
+            ),
+            4,
+            [3, 0, 2, 1],
+        ),
+    ],
+)
+def test_parse_logits(logits, length, expected_order):
+    output = parse_logits(logits, length=length)
+    assert output == expected_order
+
+
+def test_load_sort_dataset(dataset):
+    train_dataset = load_dataset_split(dataset, "train")
+    print(train_dataset)
+
+    assert len(train_dataset) == 2
+    assert sorted(train_dataset.column_names) == [
+        "sample_id",
+        "separators",
+        "source_boxes",
+        "source_classes",
+        "target_boxes",
+        "target_classes",
+        "target_index",
+    ]
+
+    sample = train_dataset[0]
+    assert sample["sample_id"] == "84b4fb2c-d62a-4e50-96f8-b0bb04410182"
+    assert sample["separators"][0] == [674, 619, 877, 620]
+    assert sample["source_boxes"][0] == [30, 141, 182, 175]
+    assert sample["target_boxes"][0] == [178, 59, 803, 87]
+    assert sample["source_classes"][0] == 11
+    assert sample["target_classes"][0] == 10
+    assert sample["target_index"][0] == 42
+
+
+@pytest.mark.parametrize(
+    ("sort_ratio", "sort_method", "expected_sorted_boxes"),
+    [
+        (
+            1,
+            "sortxy",
+            [
+                [0, 150, 100, 300],
+                [1, 0, 100, 100],
+                [200, 255, 400, 400],
+                [202, 0, 400, 250],
+            ],
+        ),
+        (
+            1,
+            "sortyx",
+            [
+                [1, 0, 100, 100],
+                [202, 0, 400, 250],
+                [0, 150, 100, 300],
+                [200, 255, 400, 400],
+            ],
+        ),
+        (
+            1,
+            "sortxy_by_column",
+            [
+                [1, 0, 100, 100],
+                [0, 150, 100, 300],
+                [202, 0, 400, 250],
+                [200, 255, 400, 400],
+            ],
+        ),
+    ],
+)
+def test_sort_dataset(sort_ratio, sort_method, expected_sorted_boxes):
+    sample = {
+        "target_classes": [
+            0,
+            0,
+            0,
+            0,
+        ],
+        "target_boxes": [
+            [1, 0, 100, 100],
+            [0, 150, 100, 300],
+            [202, 0, 400, 250],
+            [200, 255, 400, 400],
+        ],
+        "target_index": [1, 2, 3, 4],
+    }
+    sorted_boxes = sort_sample(
+        sample,
+        sort_ratio=sort_ratio,
+        sort_method=sort_method,
+    )["source_boxes"]
+    assert sorted_boxes == expected_sorted_boxes

tests/test_predict.py

@@ -0,0 +1,64 @@
+import json
+
+import pytest
+
+from layout_reader.helpers import load_model
+from layout_reader.inference import predict, run
+from tests import FIXTURES
+
+
+@pytest.fixture
+def model() -> str:
+    return str(FIXTURES / "model")
+
+
+@pytest.fixture
+def lr_dataset() -> str:
+    return str(FIXTURES / "lr_dataset")
+
+
+@pytest.fixture
+def split() -> str:
+    return "train"
+
+
+@pytest.fixture
+def images() -> str:
+    return str(FIXTURES / "lr_dataset")
+
+
+@pytest.fixture
+def expected_predictions() -> dict:
+    return json.loads((FIXTURES / "predictions.json").read_text())
+
+
+@pytest.mark.parametrize(
+    ("boxes", "classes", "separators", "expected_order"),
+    [
+        ([], [], [], []),
+        ([[1, 1, 2, 2]], [], [], [0]),
+        ([[1, 1, 2, 2], [2, 2, 3, 3]], [10, 11], [[1, 1, 1, 6]], [0, 1]),
+        ([[1, 1, 2, 2], [2, 2, 3, 3]], [], [[1, 1, 1, 6]], [0, 1]),
+        ([[1, 1, 2, 2], [2, 2, 3, 3]], [10, 11], [], [0, 1]),
+        ([[1, 1, 2, 2], [2, 2, 3, 3]], [], [], [0, 1]),
+    ],
+)
+def test_predict(model, boxes, classes, separators, expected_order):
+    model = load_model(model)
+    predicted_order = predict(model, boxes, classes, separators)
+    assert predicted_order == expected_order
+
+
+def test_run_inference(lr_dataset, split, model, tmp_path, expected_predictions):
+    output_dir = tmp_path / "output"
+
+    run(
+        dataset=lr_dataset,
+        split=split,
+        model=model,
+        output_dir=output_dir,
+    )
+
+    assert (output_dir / "predictions.json").exists()
+    predictions = json.loads((output_dir / "predictions.json").read_text())
+    assert predictions == expected_predictions