diff --git a/doctr/utils/visualization.py b/doctr/utils/visualization.py index 7f67c6e7f3..db2c97d255 100644 --- a/doctr/utils/visualization.py +++ b/doctr/utils/visualization.py @@ -42,7 +42,7 @@ def rect_patch( Returns: a rectangular Patch """ - if len(geometry) != 2 or any(not isinstance(elt, tuple) or len(elt) != 2 for elt in geometry): + if len(geometry) != 2 or any(not isinstance(elt, (tuple, list)) or len(elt) != 2 for elt in geometry): raise ValueError("invalid geometry format") # Unpack @@ -124,7 +124,8 @@ def create_obj_patch( Returns: a matplotlib Patch """ - if isinstance(geometry, tuple): + # A page export read back from a JSON file stores the geometries as lists instead of tuples + if isinstance(geometry, (tuple, list)): if len(geometry) == 2: # straight word BB (2 pts) return rect_patch(geometry, page_dimensions, **kwargs) elif len(geometry) == 4: # rotated word BB (4 pts) diff --git a/tests/common/test_io_elements.py b/tests/common/test_io_elements.py index 15aa71ae36..5f0507bc06 100644 --- a/tests/common/test_io_elements.py +++ b/tests/common/test_io_elements.py @@ -689,6 +689,8 @@ def test_page_from_dict_restores_the_image(): restored = elements.Page.from_dict(exported, page=page.page) assert np.array_equal(restored.page, page.page) assert elements.Page.from_dict(json.loads(json.dumps(exported))).render() == restored.render() + # A page restored from a JSON file can still be displayed + elements.Page.from_dict(json.loads(json.dumps(exported)), page=page.page).show(block=False) def test_kie_page_and_kie_document_from_dict(): diff --git a/tests/common/test_utils_visualization.py b/tests/common/test_utils_visualization.py index 61187d9224..124c10520d 100644 --- a/tests/common/test_utils_visualization.py +++ b/tests/common/test_utils_visualization.py @@ -1,3 +1,5 @@ +import json + import numpy as np import pytest from test_io_elements import _mock_kie_pages, _mock_layout, _mock_pages @@ -26,6 +28,28 @@ def test_visualize_page(): visualization.create_obj_patch((1, 2, 3, 4, 5), (100, 100)) +def test_visualize_json_exports(): + # A page export read back from a JSON file stores the geometries as lists instead of tuples + image = np.ones((300, 200, 3)) + page_export = _mock_pages()[0].export() + page_export["layout"] = [region.export() for region in _mock_layout()] + page_export["blocks"][0]["lines"][0]["words"][0]["geometry"] = ((0.1, 0.1), (0.4, 0.1), (0.4, 0.2), (0.1, 0.2)) + json_export = json.loads(json.dumps(page_export)) + for interactive in (True, False): + visualization.visualize_page(json_export, image, words_only=False, interactive=interactive) + fig = visualization.visualize_page(page_export, image, words_only=False, interactive=interactive) + json_fig = visualization.visualize_page(json_export, image, words_only=False, interactive=interactive) + assert len(json_fig.axes[0].patches) == len(fig.axes[0].patches) + assert [text.get_text() for text in json_fig.axes[0].texts] == [text.get_text() for text in fig.axes[0].texts] + + kie_export = json.loads(json.dumps(_mock_kie_pages()[0].export())) + fig = visualization.visualize_kie_page(kie_export, image, interactive=False, add_labels=True) + assert [text.get_text() for text in fig.axes[0].texts] == ["hello", "world"] + + with pytest.raises(ValueError): + visualization.create_obj_patch([[0.1, 0.1], [0.2]], (100, 100)) + + def test_visualize_kie_page(): pages = _mock_kie_pages() image = np.ones((300, 200, 3))