From 83dd92726ef2d0d93e534e1f38409a1f5bd9aadf Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 17 Jul 2024 09:42:37 +0200 Subject: [PATCH] fix `from_matplotlib` when Matplotlib colormaps does not provide colors property --- supervision/draw/color.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/supervision/draw/color.py b/supervision/draw/color.py index f3c1474a..7b481875 100644 --- a/supervision/draw/color.py +++ b/supervision/draw/color.py @@ -352,11 +352,17 @@ class ColorPalette: supervision-annotator-examples/visualized_color_palette.png) """ # noqa: E501 // docs mpl_palette = plt.get_cmap(palette_name, color_count) - colors = [ + + if hasattr(mpl_palette, 'colors'): + colors = mpl_palette.colors + else: + colors = [mpl_palette(i / (color_count - 1)) for i in range(color_count)] + + return cls([ Color(int(r * 255), int(g * 255), int(b * 255)) - for r, g, b, _ in mpl_palette.colors - ] - return cls(colors) + for r, g, b, _ + in colors + ]) def by_idx(self, idx: int) -> Color: """