From f96c5db223ad339121cd497721b91624629fafa7 Mon Sep 17 00:00:00 2001 From: Tai An Date: Thu, 13 Aug 2026 03:15:48 -0700 Subject: [PATCH] fix(visualization): honor add_labels in visualize_kie_page visualize_kie_page() accepts and documents add_labels ("for static plot, adds text labels on top of bounding box") but never drew any text, so a non-interactive KIE figure came out as unlabelled coloured boxes and the argument was silently a no-op. visualize_page() already draws the label in the same place, on the elif branch after the interactive one. Mirror that branch for predictions, using the per-key colour and drawing only for straight (two-point) geometries, as visualize_page() does. --- doctr/utils/visualization.py | 10 ++++++++++ tests/common/test_utils_visualization.py | 6 ++++++ 2 files changed, 16 insertions(+) diff --git a/doctr/utils/visualization.py b/doctr/utils/visualization.py index ce68b01701..7f67c6e7f3 100644 --- a/doctr/utils/visualization.py +++ b/doctr/utils/visualization.py @@ -388,6 +388,16 @@ def visualize_kie_page( if interactive: # add patch to cursor's artists artists.append(rect) + elif add_labels and len(prediction["geometry"]) == 2: + # We draw only if boxes are in straight format + ax.text( + int(page["dimensions"][1] * prediction["geometry"][0][0]), + int(page["dimensions"][0] * prediction["geometry"][0][1]), + prediction["value"], + size=10, + alpha=0.5, + color=colors[key], + ) if interactive: import mplcursors diff --git a/tests/common/test_utils_visualization.py b/tests/common/test_utils_visualization.py index cc003c416b..61187d9224 100644 --- a/tests/common/test_utils_visualization.py +++ b/tests/common/test_utils_visualization.py @@ -38,6 +38,12 @@ def test_visualize_kie_page(): visualization.visualize_kie_page(page_export, image, words_only=False, display_layout=True) visualization.visualize_kie_page(page_export, image, words_only=False, display_layout=False, interactive=False) + # static plots label each prediction when add_labels is set + fig = visualization.visualize_kie_page(pages[0].export(), image, interactive=False, add_labels=True) + assert [text.get_text() for text in fig.axes[0].texts] == ["hello", "world"] + fig = visualization.visualize_kie_page(pages[0].export(), image, interactive=False, add_labels=False) + assert len(fig.axes[0].texts) == 0 + def test_draw_boxes(): image = np.ones((256, 256, 3), dtype=np.float32)