From fe09ef9d46b46a23c21cdcee4f438bcae0b39cfd Mon Sep 17 00:00:00 2001 From: RavSinghChandan Date: Thu, 29 Oct 2026 19:41:31 +0530 Subject: [PATCH] [utils] Return exactly num_colors colors from get_colors --- doctr/utils/visualization.py | 5 +++-- tests/common/test_utils_visualization.py | 7 +++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/doctr/utils/visualization.py b/doctr/utils/visualization.py index db2c97d255..4f4678e140 100644 --- a/doctr/utils/visualization.py +++ b/doctr/utils/visualization.py @@ -145,8 +145,9 @@ def get_colors(num_colors: int) -> list[tuple[float, float, float]]: colors: list of generated colors """ colors = [] - for i in np.arange(0.0, 360.0, 360.0 / num_colors): - hue = i / 360.0 + # Integer steps give exactly num_colors hues (a float step can add one more, close to 360°, i.e. red again) + for i in range(num_colors): + hue = i / num_colors lightness = (50 + np.random.rand() * 10) / 100.0 saturation = (90 + np.random.rand() * 10) / 100.0 colors.append(colorsys.hls_to_rgb(hue, lightness, saturation)) diff --git a/tests/common/test_utils_visualization.py b/tests/common/test_utils_visualization.py index 124c10520d..7ab4c381bf 100644 --- a/tests/common/test_utils_visualization.py +++ b/tests/common/test_utils_visualization.py @@ -69,6 +69,13 @@ def test_visualize_kie_page(): assert len(fig.axes[0].texts) == 0 +@pytest.mark.parametrize("num_colors", [0, 1, 3, 161, 175]) +def test_get_colors(num_colors): + colors = visualization.get_colors(num_colors) + assert len(colors) == num_colors + assert all(len(color) == 3 and all(0 <= c <= 1 for c in color) for color in colors) + + def test_draw_boxes(): image = np.ones((256, 256, 3), dtype=np.float32) boxes = [