From 05ab96984a7f22cafc0ea63105627244b9bdb8d6 Mon Sep 17 00:00:00 2001 From: quotentiroler Date: Thu, 22 Oct 2026 20:00:46 +0200 Subject: [PATCH] [tests] Compare the table compile test's outputs on a real table image test_torch_compiled_models says "Compare that outputs are close" and then asserts only that both predictors returned one page and that the compiled one carries the expected keys. Its input is random noise, on which the predictor finds no cells, so there is nothing to compare even in principle and a compile regression would not show. Take the table image from the v1.1.0 release through a mock_table_image fixture, and compare the two predictions: the number of cells, the table's size, and every cell's geometry and score at atol=1e-4, with the logical row and column indices exact. Asserting that a cell was found at all is what keeps the comparison from passing over an empty list. --- tests/conftest.py | 11 +++++++++++ .../pytorch/test_models_table_structure_pt.py | 19 ++++++++++++++----- 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index c1a785e87a..43636b8b44 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -53,6 +53,17 @@ def mock_payslip(tmpdir_factory): return fn +@pytest.fixture(scope="session") +def mock_table_image(tmpdir_factory): + url = "https://github.com/mindee/doctr/releases/download/v1.1.0/table_structure_example.png" + file = BytesIO(requests.get(url).content) + folder = tmpdir_factory.mktemp("data") + fn = str(folder.join("mock_table.png")) + with open(fn, "wb") as f: + f.write(file.getbuffer()) + return fn + + @pytest.fixture(scope="session") def mock_figure_page(mock_image_stream, tmpdir_factory): # A heading, two paragraphs and a captioned photograph in between diff --git a/tests/pytorch/test_models_table_structure_pt.py b/tests/pytorch/test_models_table_structure_pt.py index e742f7519e..a12e9dd07b 100644 --- a/tests/pytorch/test_models_table_structure_pt.py +++ b/tests/pytorch/test_models_table_structure_pt.py @@ -6,6 +6,7 @@ import pytest import torch +from doctr.io import DocumentFile from doctr.models import table_structure from doctr.models.table_structure import TableCenterNet from doctr.models.table_structure.predictor import TablePredictor @@ -161,19 +162,27 @@ def test_models_onnx_export(arch_name, input_shape): "tablecenternet", ], ) -def test_torch_compiled_models(arch_name): - page = (255 * np.random.rand(1024, 1024, 3)).astype(np.uint8) +def test_torch_compiled_models(arch_name, mock_table_image): + doc = DocumentFile.from_images([mock_table_image]) predictor = table_structure.zoo.table_predictor(arch_name, pretrained=True) assert isinstance(predictor, TablePredictor) - out = predictor([page]) + out = predictor(doc) # Compile the model base = table_structure.__dict__[arch_name](pretrained=True).eval() compiled_model = torch.compile(base) assert isinstance(compiled_model, _CompiledModule) compiled_predictor = table_structure.zoo.table_predictor(compiled_model) - compiled_out = compiled_predictor([page]) + compiled_out = compiled_predictor(doc) # Compare that outputs are close assert len(out) == len(compiled_out) == 1 - assert {"cells", "num_rows", "num_cols"} <= set(compiled_out[0]) + table, compiled_table = out[0], compiled_out[0] + # A table was actually found, otherwise there is nothing to compare + assert len(table["cells"]) > 0 + assert len(table["cells"]) == len(compiled_table["cells"]) + assert (table["num_rows"], table["num_cols"]) == (compiled_table["num_rows"], compiled_table["num_cols"]) + for cell, compiled_cell in zip(table["cells"], compiled_table["cells"]): + assert np.allclose(cell["geometry"], compiled_cell["geometry"], atol=1e-4) + assert np.isclose(cell["score"], compiled_cell["score"], atol=1e-4) + assert all(cell[k] == compiled_cell[k] for k in ("row_start", "row_end", "col_start", "col_end"))