diff --git a/doctr/io/elements.py b/doctr/io/elements.py index 54fa7550c9..81b8f550cc 100644 --- a/doctr/io/elements.py +++ b/doctr/io/elements.py @@ -408,9 +408,8 @@ def __init__( if geometry is None: line_boxes = [word.geometry for line in lines for word in line.words] artefact_boxes = [artefact.geometry for artefact in artefacts] - box_resolution_fn = ( - resolve_enclosing_rbbox if isinstance(lines[0].geometry, np.ndarray) else resolve_enclosing_bbox - ) + # Check whether this is a rotated or straight box (geometries read back from an export are lists) + box_resolution_fn = resolve_enclosing_rbbox if len(lines[0].geometry) == 4 else resolve_enclosing_bbox geometry = box_resolution_fn(line_boxes + artefact_boxes) # type: ignore super().__init__(lines=lines, artefacts=artefacts) diff --git a/tests/common/test_io_elements.py b/tests/common/test_io_elements.py index 5f0507bc06..5fbfcc94b1 100644 --- a/tests/common/test_io_elements.py +++ b/tests/common/test_io_elements.py @@ -442,6 +442,25 @@ def test_block(): } +def test_block_rotated_lines_from_dict(): + # Rotated geometries read back from an export are nested lists, not numpy arrays + polygon = [[0.1, 0.1], [0.4, 0.2], [0.35, 0.4], [0.05, 0.3]] + word = elements.Word("hello", 0.9, polygon, 0.9, {"value": 0, "confidence": None}) + line = elements.Line.from_dict(elements.Line([word], geometry=polygon).export()) + assert isinstance(line.geometry, list) + + # The block geometry is resolved as a rotated box, not as a straight one + block = elements.Block([line]) + assert np.asarray(block.geometry).shape == (4, 2) + + # Exporting a page loaded from a dict keeps the rotated block geometry + page = elements.Page(np.zeros((10, 10, 3), dtype=np.uint8), [elements.Block([line])], 0, (10, 10)) + exported = json.loads(json.dumps(page.export())) + reloaded = elements.Page.from_dict(exported).export() + assert np.asarray(reloaded["blocks"][0]["geometry"]).shape == (4, 2) + np.testing.assert_allclose(reloaded["blocks"][0]["geometry"], exported["blocks"][0]["geometry"], atol=1e-6) + + def test_page(): page = np.zeros((300, 200, 3), dtype=np.uint8) page_idx = 0