feat: visualization custom label colours

model_api/examples/visualization/README.md

@@ -29,3 +29,22 @@ and then run
 ```bash
 uv run python examples/visualization/run.py --image data/cards.png --model data/otx_models/ssd-card-detection.xml --output cards_result.jpg
 ```
+
+## Matching your own label colours
+
+By default the `Visualizer` assigns a colour to each label from a built-in palette. Pass
+`label_colors` to render predictions with colours you control, for example the label
+colours defined in your project:
+
+```python
+from model_api.visualizer import Visualizer
+
+visualizer = Visualizer(label_colors={"car": "#FF0000", "person": (0, 255, 0)})
+visualizer.show(image, result)
+```
+
+Colours are either any string accepted by PIL (`"#RRGGBB"`, `"red"`, ...) or an
+`(R, G, B)` tuple of integers in the 0-255 range. Labels that are not in the mapping keep
+their default colour, and invalid colours raise a `ValueError` when the `Visualizer` is
+created. The mapping is applied to detection, instance segmentation, classification and
+anomaly results.

model_api/src/model_api/visualizer/defaults.py

@@ -23,3 +23,10 @@
 SCALE_BASELINE: int = 1280
 """Longer-edge pixel count of 720p (landscape). Used as the denominator when
 computing the auto-scale factor."""
+
+# Colors
+DEFAULT_SHAPE_COLOR: str = "blue"
+"""Default color for bounding boxes and polygons."""
+
+DEFAULT_LABEL_BG_COLOR: str = "yellow"
+"""Default background color for label chips."""

model_api/src/model_api/visualizer/primitive/bounding_box.py

@@ -7,7 +7,7 @@
 
 from PIL import Image, ImageDraw
 
-from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_OUTLINE_WIDTH
+from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_OUTLINE_WIDTH, DEFAULT_SHAPE_COLOR
 from model_api.visualizer.utils import default_font, make_label_image
 
 from .primitive import Primitive
@@ -38,7 +38,7 @@ def __init__(
         x2: int,
         y2: int,
         label: str | None = None,
-        color: str | tuple[int, int, int] = "blue",
+        color: str | tuple[int, int, int] = DEFAULT_SHAPE_COLOR,
         outline_width: int = DEFAULT_OUTLINE_WIDTH,
         font_size: int = DEFAULT_FONT_SIZE,
     ) -> None:

model_api/src/model_api/visualizer/primitive/label.py

@@ -8,7 +8,7 @@
 
 from PIL import Image
 
-from model_api.visualizer.defaults import DEFAULT_FONT_SIZE
+from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_LABEL_BG_COLOR
 from model_api.visualizer.utils import default_font, make_label_image, truetype_font
 
 from .primitive import Primitive
@@ -47,7 +47,7 @@ def __init__(
         label: str,
         score: Union[float, None] = None,
         fg_color: Union[str, tuple[int, int, int]] = "black",
-        bg_color: Union[str, tuple[int, int, int]] = "yellow",
+        bg_color: Union[str, tuple[int, int, int]] = DEFAULT_LABEL_BG_COLOR,
         font_path: Union[str, BytesIO, None] = None,
         size: int = DEFAULT_FONT_SIZE,
     ) -> None:

model_api/src/model_api/visualizer/primitive/polygon.py

@@ -9,9 +9,10 @@
 from typing import TYPE_CHECKING
 
 import cv2
-from PIL import Image, ImageColor, ImageDraw
+from PIL import Image, ImageDraw
 
-from model_api.visualizer.defaults import DEFAULT_OPACITY, DEFAULT_OUTLINE_WIDTH
+from model_api.visualizer.defaults import DEFAULT_OPACITY, DEFAULT_OUTLINE_WIDTH, DEFAULT_SHAPE_COLOR
+from model_api.visualizer.utils import to_rgb
 
 from .primitive import Primitive
 
@@ -42,7 +43,7 @@ def __init__(
         self,
         points: list[tuple[int, int]] | None = None,
         mask: np.ndarray | None = None,
-        color: str | tuple[int, int, int] = "blue",
+        color: str | tuple[int, int, int] = DEFAULT_SHAPE_COLOR,
         opacity: float = DEFAULT_OPACITY,
         outline_width: int = DEFAULT_OUTLINE_WIDTH,
     ) -> None:
@@ -112,6 +113,6 @@ def compute(self, image: Image) -> Image:
 
         draw = ImageDraw.Draw(image, "RGBA")
         # Draw polygon with darker edge and a semi-transparent fill.
-        ink = ImageColor.getrgb(self.color)
+        ink = to_rgb(self.color)
         draw.polygon(self.points, fill=(*ink, int(255 * self.opacity)), outline=self.color, width=self.outline_width)
         return image

model_api/src/model_api/visualizer/scene/anomaly.py

@@ -3,30 +3,55 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 import cv2
 from PIL import Image
 
 from model_api.models.result import AnomalyResult
-from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_OUTLINE_WIDTH
+from model_api.visualizer.defaults import (
+    DEFAULT_FONT_SIZE,
+    DEFAULT_LABEL_BG_COLOR,
+    DEFAULT_OUTLINE_WIDTH,
+    DEFAULT_SHAPE_COLOR,
+)
 from model_api.visualizer.layout import Flatten, Layout
 from model_api.visualizer.primitive import BoundingBox, Label, Overlay, Polygon
 
 from .scene import Scene
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from model_api.visualizer.utils import Color
+
 
 class AnomalyScene(Scene):
-    """Anomaly Scene."""
+    """Anomaly Scene.
+
+    Args:
+        image: Base image to draw on.
+        result: Anomaly result to render.
+        layout: Optional layout to use for rendering.
+        scale: Scale factor applied to drawing sizes.
+        label_colors: Optional mapping of label name to colour. When the predicted
+            label is present in the mapping, its colour is used for the label
+            background, the bounding boxes and the mask polygon. Otherwise the
+            default colours are used.
+    """
 
     def __init__(
         self,
         image: Image,
         result: AnomalyResult,
         layout: Union[Layout, None] = None,
         scale: float = 1.0,
+        label_colors: Union["Mapping[str, Color]", None] = None,
     ) -> None:
         self.scale = scale
+        color = (label_colors or {}).get(result.pred_label) if result.pred_label is not None else None
+        self.shape_color: Color = DEFAULT_SHAPE_COLOR if color is None else color
+        self.label_bg_color: Color = DEFAULT_LABEL_BG_COLOR if color is None else color
         super().__init__(
             base=image,
             overlay=self._get_overlays(result),
@@ -50,6 +75,7 @@ def _get_bounding_boxes(self, result: AnomalyResult) -> list[BoundingBox]:
                     y1=box[1],
                     x2=box[2],
                     y2=box[3],
+                    color=self.shape_color,
                     outline_width=max(1, int(DEFAULT_OUTLINE_WIDTH * self.scale)),
                     font_size=int(DEFAULT_FONT_SIZE * self.scale),
                 )
@@ -64,6 +90,7 @@ def _get_labels(self, result: AnomalyResult) -> list[Label]:
                 Label(
                     label=result.pred_label,
                     score=result.pred_score,
+                    bg_color=self.label_bg_color,
                     size=int(DEFAULT_FONT_SIZE * self.scale),
                 ),
             )
@@ -74,6 +101,7 @@ def _get_polygons(self, result: AnomalyResult) -> list[Polygon]:
             return [
                 Polygon(
                     result.pred_mask,
+                    color=self.shape_color,
                     outline_width=max(1, int(DEFAULT_OUTLINE_WIDTH * self.scale)),
                 ),
             ]

model_api/src/model_api/visualizer/scene/classification.py

@@ -3,30 +3,46 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 import cv2
 from PIL import Image
 
 from model_api.models.result import ClassificationResult
-from model_api.visualizer.defaults import DEFAULT_FONT_SIZE
+from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_LABEL_BG_COLOR
 from model_api.visualizer.layout import Flatten, Layout
 from model_api.visualizer.primitive import Label, Overlay
 
 from .scene import Scene
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from model_api.visualizer.utils import Color
+
 
 class ClassificationScene(Scene):
-    """Classification Scene."""
+    """Classification Scene.
+
+    Args:
+        image: Base image to draw on.
+        result: Classification result to render.
+        layout: Optional layout to use for rendering.
+        scale: Scale factor applied to drawing sizes.
+        label_colors: Optional mapping of label name to colour used as the label
+            background. Labels absent from the mapping keep the default background.
+    """
 
     def __init__(
         self,
         image: Image,
         result: ClassificationResult,
         layout: Union[Layout, None] = None,
         scale: float = 1.0,
+        label_colors: Union["Mapping[str, Color]", None] = None,
     ) -> None:
         self.scale = scale
+        self.label_colors = label_colors or {}
         super().__init__(
             base=image,
             label=self._get_labels(result),
@@ -44,6 +60,7 @@ def _get_labels(self, result: ClassificationResult) -> list[Label]:
                             label=label.name,
                             score=label.confidence,
                             size=int(DEFAULT_FONT_SIZE * self.scale),
+                            bg_color=self.label_colors.get(label.name, DEFAULT_LABEL_BG_COLOR),
                         ),
                     )
         return labels

model_api/src/model_api/visualizer/scene/detection.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 import cv2
 from PIL import Image
@@ -16,18 +16,33 @@
 
 from .scene import Scene
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from model_api.visualizer.utils import Color
+
 
 class DetectionScene(Scene):
-    """Detection Scene."""
+    """Detection Scene.
+
+    Args:
+        image: Base image to draw on.
+        result: Detection result to render.
+        layout: Optional layout to use for rendering.
+        scale: Scale factor applied to drawing sizes.
+        label_colors: Optional mapping of label name to colour. Labels absent from the
+            mapping keep their automatically assigned palette colour.
+    """
 
     def __init__(
         self,
         image: Image,
         result: DetectionResult,
         layout: Union[Layout, None] = None,
         scale: float = 1.0,
+        label_colors: Union["Mapping[str, Color]", None] = None,
     ) -> None:
-        self.color_per_label = get_label_color_mapping(result.label_names)
+        self.color_per_label = get_label_color_mapping(result.label_names, overrides=label_colors)
         self.scale = scale
         super().__init__(
             base=image,

model_api/src/model_api/visualizer/scene/segmentation/instance_segmentation.py

@@ -3,7 +3,7 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 import cv2
 from PIL import Image
@@ -15,18 +15,33 @@
 from model_api.visualizer.scene import Scene
 from model_api.visualizer.utils import get_label_color_mapping
 
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from model_api.visualizer.utils import Color
+
 
 class InstanceSegmentationScene(Scene):
-    """Instance Segmentation Scene."""
+    """Instance Segmentation Scene.
+
+    Args:
+        image: Base image to draw on.
+        result: Instance segmentation result to render.
+        layout: Optional layout to use for rendering.
+        scale: Scale factor applied to drawing sizes.
+        label_colors: Optional mapping of label name to colour. Labels absent from the
+            mapping keep their automatically assigned palette colour.
+    """
 
     def __init__(
         self,
         image: Image,
         result: InstanceSegmentationResult,
         layout: Union[Layout, None] = None,
         scale: float = 1.0,
+        label_colors: Union["Mapping[str, Color]", None] = None,
     ) -> None:
-        self.color_per_label = get_label_color_mapping(result.label_names)
+        self.color_per_label = get_label_color_mapping(result.label_names, overrides=label_colors)
         self.scale = scale
         super().__init__(
             base=image,