diff --git a/model_api/examples/visualization/README.md b/model_api/examples/visualization/README.md index 565dc4861..04b7a5593 100644 --- a/model_api/examples/visualization/README.md +++ b/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. diff --git a/model_api/src/model_api/visualizer/defaults.py b/model_api/src/model_api/visualizer/defaults.py index 0411658c1..796947339 100644 --- a/model_api/src/model_api/visualizer/defaults.py +++ b/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.""" diff --git a/model_api/src/model_api/visualizer/primitive/bounding_box.py b/model_api/src/model_api/visualizer/primitive/bounding_box.py index c5dc27ead..94d15318b 100644 --- a/model_api/src/model_api/visualizer/primitive/bounding_box.py +++ b/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: diff --git a/model_api/src/model_api/visualizer/primitive/label.py b/model_api/src/model_api/visualizer/primitive/label.py index fb72bceca..874f52a75 100644 --- a/model_api/src/model_api/visualizer/primitive/label.py +++ b/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: diff --git a/model_api/src/model_api/visualizer/primitive/polygon.py b/model_api/src/model_api/visualizer/primitive/polygon.py index 00fd95350..fc6a7ab87 100644 --- a/model_api/src/model_api/visualizer/primitive/polygon.py +++ b/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 diff --git a/model_api/src/model_api/visualizer/scene/anomaly.py b/model_api/src/model_api/visualizer/scene/anomaly.py index 9c196fbc9..dfcae11e6 100644 --- a/model_api/src/model_api/visualizer/scene/anomaly.py +++ b/model_api/src/model_api/visualizer/scene/anomaly.py @@ -3,21 +3,42 @@ # 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, @@ -25,8 +46,12 @@ def __init__( 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)), ), ] diff --git a/model_api/src/model_api/visualizer/scene/classification.py b/model_api/src/model_api/visualizer/scene/classification.py index 94160f966..821af20f6 100644 --- a/model_api/src/model_api/visualizer/scene/classification.py +++ b/model_api/src/model_api/visualizer/scene/classification.py @@ -3,21 +3,35 @@ # 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, @@ -25,8 +39,10 @@ def __init__( 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 diff --git a/model_api/src/model_api/visualizer/scene/detection.py b/model_api/src/model_api/visualizer/scene/detection.py index fd2a98b55..eeced635b 100644 --- a/model_api/src/model_api/visualizer/scene/detection.py +++ b/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,9 +16,23 @@ 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, @@ -26,8 +40,9 @@ def __init__( 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, diff --git a/model_api/src/model_api/visualizer/scene/segmentation/instance_segmentation.py b/model_api/src/model_api/visualizer/scene/segmentation/instance_segmentation.py index 3c0044b45..650326aa4 100644 --- a/model_api/src/model_api/visualizer/scene/segmentation/instance_segmentation.py +++ b/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,9 +15,23 @@ 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, @@ -25,8 +39,9 @@ def __init__( 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, diff --git a/model_api/src/model_api/visualizer/utils.py b/model_api/src/model_api/visualizer/utils.py index c982be5f3..6f3a7031c 100644 --- a/model_api/src/model_api/visualizer/utils.py +++ b/model_api/src/model_api/visualizer/utils.py @@ -3,13 +3,18 @@ from __future__ import annotations from functools import lru_cache -from typing import Union +from typing import TYPE_CHECKING, Union -from PIL import Image, ImageDraw, ImageFont +from PIL import Image, ImageColor, ImageDraw, ImageFont + +if TYPE_CHECKING: + from collections.abc import Mapping # Copyright (C) 2026 Intel Corporation # SPDX-License-Identifier: Apache-2.0 +Color = Union[str, tuple[int, int, int]] + COLOR_PALETTE = [ "#FF6B6B", # Red "#4ECDC4", # Teal @@ -34,17 +39,90 @@ ] -def get_label_color_mapping(labels: list[str]) -> dict[str, str]: +def _is_valid_rgb_tuple(color: tuple) -> bool: + """Check that a tuple is a valid ``(R, G, B)`` triplet. + + Args: + color: Tuple to check. + + Returns: + True when the tuple holds exactly three integers in the 0-255 range. + """ + return len(color) == 3 and all(isinstance(channel, int) and 0 <= channel <= 255 for channel in color) + + +def to_rgb(color: Color) -> tuple[int, int, int]: + """Normalize a colour to an ``(R, G, B)`` tuple. + + ``PIL.ImageColor.getrgb`` only accepts strings, so tuples are returned unchanged. + + Args: + color: Colour string accepted by PIL (e.g. ``"#RRGGBB"`` or ``"red"``) or an + ``(R, G, B)`` tuple of integers in the 0-255 range. + + Returns: + The colour as an ``(R, G, B)`` tuple. + """ + if isinstance(color, str): + return ImageColor.getrgb(color) + return color + + +def validate_label_colors(label_colors: Union[Mapping[str, Color], None]) -> dict[str, Color]: + """Validate a mapping of label name to colour. + + Args: + label_colors: Mapping of label name to a colour. Colours are either a string + accepted by PIL (e.g. ``"#RRGGBB"`` or ``"red"``) or an ``(R, G, B)`` tuple + of integers in the 0-255 range. ``None`` is treated as an empty mapping. + + Returns: + A copy of the mapping as a plain dictionary. Empty when *label_colors* is None. + + Raises: + ValueError: If any colour is not a valid colour string or RGB tuple. + """ + if label_colors is None: + return {} + + validated: dict[str, Color] = {} + for label, color in label_colors.items(): + if isinstance(color, str): + try: + to_rgb(color) + except ValueError as error: + msg = f"Invalid color {color!r} for label {label!r}." + raise ValueError(msg) from error + elif not (isinstance(color, tuple) and _is_valid_rgb_tuple(color)): + msg = ( + f"Invalid color {color!r} for label {label!r}. Expected a color string or " + "a tuple of three integers in the 0-255 range." + ) + raise ValueError(msg) + validated[label] = color + return validated + + +def get_label_color_mapping( + labels: list[str], + overrides: Union[Mapping[str, Color], None] = None, +) -> dict[str, Color]: """Generate a consistent color mapping for a list of labels. Args: labels: List of label names. + overrides: Optional mapping of label name to colour. Entries whose label appears + in *labels* replace the automatically assigned palette colour; other entries + are ignored. Returns: - Dictionary mapping each label to a hex color string. + Dictionary mapping each label to a colour. """ unique_labels = sorted(set(labels)) - return {label: COLOR_PALETTE[i % len(COLOR_PALETTE)] for i, label in enumerate(unique_labels)} + mapping: dict[str, Color] = {label: COLOR_PALETTE[i % len(COLOR_PALETTE)] for i, label in enumerate(unique_labels)} + if overrides: + mapping.update({label: color for label, color in overrides.items() if label in mapping}) + return mapping @lru_cache(maxsize=5) diff --git a/model_api/src/model_api/visualizer/visualizer.py b/model_api/src/model_api/visualizer/visualizer.py index 33609bf8b..321f94511 100644 --- a/model_api/src/model_api/visualizer/visualizer.py +++ b/model_api/src/model_api/visualizer/visualizer.py @@ -30,11 +30,14 @@ Scene, SegmentationScene, ) +from .utils import validate_label_colors if TYPE_CHECKING: + from collections.abc import Mapping from pathlib import Path from .layout import Layout + from .utils import Color class Visualizer: @@ -45,11 +48,30 @@ class Visualizer: auto_scale: When True, drawing sizes (line widths, font sizes, etc.) are automatically scaled relative to 720p so that annotations remain visible on high-resolution images. Defaults to True. + label_colors: Optional mapping of label name to colour, used to render + predictions with colours defined by the caller (for example the label + colours of a project) instead of the automatically assigned palette. + Colours are either a string accepted by PIL (e.g. ``"#RRGGBB"`` or + ``"red"``) or an ``(R, G, B)`` tuple of integers in the 0-255 range. + Labels absent from the mapping keep their default colour. + + Raises: + ValueError: If *label_colors* contains an invalid colour. + + Example: + >>> visualizer = Visualizer(label_colors={"car": "#FF0000", "person": (0, 255, 0)}) + >>> visualizer.show(image, result) """ - def __init__(self, layout: Layout | None = None, auto_scale: bool = True) -> None: + def __init__( + self, + layout: Layout | None = None, + auto_scale: bool = True, + label_colors: Mapping[str, Color] | None = None, + ) -> None: self.layout = layout self.auto_scale = auto_scale + self.label_colors = validate_label_colors(label_colors) @staticmethod def compute_scale_factor(image: Image.Image) -> float: @@ -135,17 +157,17 @@ def _scene_from_result(self, image: Image, result: Result) -> Scene: scene: Scene if isinstance(result, AnomalyResult): - scene = AnomalyScene(image, result, self.layout, scale=scale) + scene = AnomalyScene(image, result, self.layout, scale=scale, label_colors=self.label_colors) elif isinstance(result, ClassificationResult): - scene = ClassificationScene(image, result, self.layout, scale=scale) + scene = ClassificationScene(image, result, self.layout, scale=scale, label_colors=self.label_colors) elif isinstance(result, InstanceSegmentationResult): # Note: This has to be before DetectionScene because InstanceSegmentationResult is a subclass # of DetectionResult - scene = InstanceSegmentationScene(image, result, self.layout, scale=scale) + scene = InstanceSegmentationScene(image, result, self.layout, scale=scale, label_colors=self.label_colors) elif isinstance(result, ImageResultWithSoftPrediction): scene = SegmentationScene(image, result, self.layout, scale=scale) elif isinstance(result, DetectionResult): - scene = DetectionScene(image, result, self.layout, scale=scale) + scene = DetectionScene(image, result, self.layout, scale=scale, label_colors=self.label_colors) elif isinstance(result, DetectedKeypoints): scene = KeypointScene(image, result, self.layout, scale=scale) else: diff --git a/model_api/tests/unit/visualizer/test_label_colors.py b/model_api/tests/unit/visualizer/test_label_colors.py new file mode 100644 index 000000000..a49707066 --- /dev/null +++ b/model_api/tests/unit/visualizer/test_label_colors.py @@ -0,0 +1,381 @@ +"""Tests for the optional label colour mapping.""" + +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +from model_api.models.result import ( + AnomalyResult, + ClassificationResult, + DetectionResult, + InstanceSegmentationResult, +) +from model_api.models.result.classification import Label +from model_api.visualizer import Visualizer +from model_api.visualizer.primitive import Polygon +from model_api.visualizer.scene.anomaly import AnomalyScene +from model_api.visualizer.scene.classification import ClassificationScene +from model_api.visualizer.scene.detection import DetectionScene +from model_api.visualizer.scene.segmentation.instance_segmentation import InstanceSegmentationScene +from model_api.visualizer.utils import COLOR_PALETTE, get_label_color_mapping, to_rgb, validate_label_colors +from PIL import Image + + +@pytest.fixture +def detection_result() -> DetectionResult: + return DetectionResult( + bboxes=np.array([[0, 0, 64, 64], [32, 32, 96, 96]]), + labels=np.array([0, 1]), + label_names=["car", "person"], + scores=np.array([0.85, 0.75]), + saliency_map=None, + ) + + +@pytest.fixture +def instance_segmentation_result() -> InstanceSegmentationResult: + return InstanceSegmentationResult( + bboxes=np.array([[0, 0, 64, 64], [32, 32, 96, 96]]), + labels=np.array([0, 1]), + masks=np.array([ + np.ones((100, 100), dtype=np.uint8), + np.ones((100, 100), dtype=np.uint8), + ]), + scores=np.array([0.85, 0.75]), + label_names=["car", "person"], + saliency_map=None, + feature_vector=np.array([1, 2, 3]), + ) + + +@pytest.fixture +def classification_result() -> ClassificationResult: + return ClassificationResult( + top_labels=[ + Label(name="cat", confidence=0.95), + Label(name="dog", confidence=0.90), + ], + saliency_map=None, + ) + + +@pytest.fixture +def anomaly_result(mock_image: Image) -> AnomalyResult: + mask = np.zeros(mock_image.size, dtype=np.uint8) + mask[32:96, 32:96] = 255 + return AnomalyResult( + anomaly_map=None, + pred_boxes=np.array([[0, 0, 64, 64]]), + pred_label="Anomaly", + pred_mask=mask, + pred_score=0.85, + ) + + +class TestValidateLabelColors: + """Tests for validate_label_colors().""" + + def test_none_returns_empty_mapping(self): + assert validate_label_colors(None) == {} + + @pytest.mark.parametrize("color", ["#FF0000", "red", (12, 34, 56)]) + def test_accepts_valid_colors(self, color): + assert validate_label_colors({"car": color}) == {"car": color} + + def test_returns_a_copy(self): + source = {"car": "#FF0000"} + validated = validate_label_colors(source) + validated["person"] = "#00FF00" + assert source == {"car": "#FF0000"} + + @pytest.mark.parametrize( + "color", + [ + "not-a-colour", + (1, 2), + (1, 2, 3, 4), + (1, 2, 300), + (1, 2, -1), + (1, 2, "3"), + 123, + None, + ], + ) + def test_rejects_invalid_colors(self, color): + with pytest.raises(ValueError, match="car"): + validate_label_colors({"car": color}) + + +class TestGetLabelColorMapping: + """Tests for get_label_color_mapping().""" + + def test_without_overrides_uses_palette(self): + mapping = get_label_color_mapping(["person", "car"]) + assert mapping == {"car": COLOR_PALETTE[0], "person": COLOR_PALETTE[1]} + + def test_overrides_replace_palette_colors(self): + mapping = get_label_color_mapping(["person", "car"], overrides={"car": "#123456"}) + assert mapping["car"] == "#123456" + assert mapping["person"] == COLOR_PALETTE[1] + + def test_overrides_accept_rgb_tuples(self): + mapping = get_label_color_mapping(["car"], overrides={"car": (1, 2, 3)}) + assert mapping["car"] == (1, 2, 3) + + def test_unknown_override_keys_are_ignored(self): + mapping = get_label_color_mapping(["car"], overrides={"bicycle": "#123456"}) + assert mapping == {"car": COLOR_PALETTE[0]} + + def test_none_overrides_behave_like_no_overrides(self): + assert get_label_color_mapping(["car"], overrides=None) == get_label_color_mapping(["car"]) + + +class TestVisualizerLabelColors: + """Tests for the Visualizer label_colors argument.""" + + def test_defaults_to_empty_mapping(self): + assert Visualizer().label_colors == {} + + def test_validates_eagerly(self): + with pytest.raises(ValueError, match="car"): + Visualizer(label_colors={"car": "not-a-colour"}) + + def test_forwards_mapping_to_detection_scene(self, mock_image: Image, detection_result: DetectionResult): + visualizer = Visualizer(label_colors={"car": "#123456"}) + scene = visualizer._scene_from_result(mock_image, detection_result) # noqa: SLF001 + assert isinstance(scene, DetectionScene) + assert scene.color_per_label["car"] == "#123456" + + def test_renders_with_mapping(self, mock_image: Image, detection_result: DetectionResult): + visualizer = Visualizer(label_colors={"car": "#123456"}) + assert isinstance(visualizer.render(mock_image, detection_result), Image.Image) + + +class TestDetectionSceneColors: + """Tests for label colours in the detection scene.""" + + def test_mapped_label_uses_custom_color(self, mock_image: Image, detection_result: DetectionResult): + scene = DetectionScene(mock_image, detection_result, label_colors={"car": "#123456"}) + assert scene.bounding_box is not None + assert scene.color_per_label["car"] == "#123456" + assert scene.bounding_box[0].color == "#123456" + + def test_unmapped_label_keeps_palette_color(self, mock_image: Image, detection_result: DetectionResult): + scene = DetectionScene(mock_image, detection_result, label_colors={"car": "#123456"}) + assert scene.bounding_box is not None + assert scene.color_per_label["person"] == COLOR_PALETTE[1] + assert scene.bounding_box[1].color == COLOR_PALETTE[1] + + def test_without_mapping_uses_palette(self, mock_image: Image, detection_result: DetectionResult): + scene = DetectionScene(mock_image, detection_result) + assert scene.color_per_label == {"car": COLOR_PALETTE[0], "person": COLOR_PALETTE[1]} + + +class TestInstanceSegmentationSceneColors: + """Tests for label colours in the instance segmentation scene.""" + + def test_polygons_use_custom_color( + self, + mock_image: Image, + instance_segmentation_result: InstanceSegmentationResult, + ): + scene = InstanceSegmentationScene( + mock_image, + instance_segmentation_result, + label_colors={"car": "#123456"}, + ) + assert scene.polygon is not None + assert scene.polygon[0].color == "#123456" + + def test_label_chips_use_custom_color( + self, + mock_image: Image, + instance_segmentation_result: InstanceSegmentationResult, + ): + scene = InstanceSegmentationScene( + mock_image, + instance_segmentation_result, + label_colors={"car": "#123456"}, + ) + assert scene.label is not None + colors = {label.label: label.bg_color for label in scene.label} + assert colors["car"] == "#123456" + assert colors["person"] == COLOR_PALETTE[1] + + def test_bounding_boxes_use_custom_color( + self, + mock_image: Image, + instance_segmentation_result: InstanceSegmentationResult, + ): + scene = InstanceSegmentationScene( + mock_image, + instance_segmentation_result, + label_colors={"car": "#123456"}, + ) + assert scene._get_bounding_boxes(instance_segmentation_result)[0].color == "#123456" # noqa: SLF001 + + +class TestClassificationSceneColors: + """Tests for label colours in the classification scene.""" + + def test_mapped_label_uses_custom_background( + self, + mock_image: Image, + classification_result: ClassificationResult, + ): + scene = ClassificationScene(mock_image, classification_result, label_colors={"cat": "#123456"}) + assert scene.label is not None + assert scene.label[0].bg_color == "#123456" + + def test_unmapped_label_keeps_default_background( + self, + mock_image: Image, + classification_result: ClassificationResult, + ): + scene = ClassificationScene(mock_image, classification_result, label_colors={"cat": "#123456"}) + assert scene.label is not None + assert scene.label[1].bg_color == "yellow" + + def test_without_mapping_keeps_default_background( + self, + mock_image: Image, + classification_result: ClassificationResult, + ): + scene = ClassificationScene(mock_image, classification_result) + assert scene.label is not None + assert [label.bg_color for label in scene.label] == ["yellow", "yellow"] + + +class TestAnomalySceneColors: + """Tests for label colours in the anomaly scene.""" + + def test_mapped_label_colors_primitives(self, mock_image: Image, anomaly_result: AnomalyResult): + scene = AnomalyScene(mock_image, anomaly_result, label_colors={"Anomaly": "#123456"}) + assert scene.label is not None + assert scene.bounding_box is not None + assert scene.polygon is not None + assert scene.label[0].bg_color == "#123456" + assert scene.bounding_box[0].color == "#123456" + assert scene.polygon[0].color == "#123456" + + def test_unmapped_label_keeps_defaults(self, mock_image: Image, anomaly_result: AnomalyResult): + scene = AnomalyScene(mock_image, anomaly_result, label_colors={"Normal": "#123456"}) + assert scene.label is not None + assert scene.bounding_box is not None + assert scene.polygon is not None + assert scene.label[0].bg_color == "yellow" + assert scene.bounding_box[0].color == "blue" + assert scene.polygon[0].color == "blue" + + def test_without_mapping_keeps_defaults(self, mock_image: Image, anomaly_result: AnomalyResult): + scene = AnomalyScene(mock_image, anomaly_result) + assert scene.label is not None + assert scene.bounding_box is not None + assert scene.polygon is not None + assert scene.label[0].bg_color == "yellow" + assert scene.bounding_box[0].color == "blue" + assert scene.polygon[0].color == "blue" + + def test_missing_pred_label_is_ignored(self, mock_image: Image): + result = AnomalyResult( + anomaly_map=None, + pred_boxes=np.array([[0, 0, 64, 64]]), + pred_label=None, + pred_mask=np.zeros((100, 100), dtype=np.uint8), + pred_score=None, + ) + scene = AnomalyScene(mock_image, result, label_colors={"Anomaly": "#123456"}) + assert scene.bounding_box is not None + assert scene.label == [] + assert scene.bounding_box[0].color == "blue" + + +class TestRgbTupleColors: + """Tuple colours must render as well as string colours. + + Regression test: Polygon used PIL's ImageColor.getrgb(), which only accepts + strings, so RGB tuples raised "'tuple' object has no attribute 'lower'" for any + result that draws polygons (e.g. MaskRCNN instance segmentation). + """ + + def test_polygon_primitive_accepts_rgb_tuple(self, mock_image: Image): + mask = np.zeros((100, 100), dtype=np.uint8) + mask[10:90, 10:90] = 1 + polygon = Polygon(mask=mask, color=(18, 52, 86)) + assert isinstance(polygon.compute(mock_image.copy()), Image.Image) + + def test_polygon_primitive_tuple_matches_hex(self, mock_image: Image): + mask = np.zeros((100, 100), dtype=np.uint8) + mask[10:90, 10:90] = 1 + from_tuple = Polygon(mask=mask, color=(18, 52, 86)).compute(mock_image.copy()) + from_hex = Polygon(mask=mask, color="#123456").compute(mock_image.copy()) + assert from_tuple.tobytes() == from_hex.tobytes() + + def test_instance_segmentation_renders_with_rgb_tuple( + self, + mock_image: Image, + instance_segmentation_result: InstanceSegmentationResult, + ): + visualizer = Visualizer(label_colors={"car": (18, 52, 86)}) + assert isinstance(visualizer.render(mock_image, instance_segmentation_result), Image.Image) + + def test_anomaly_renders_with_rgb_tuple(self, mock_image: Image, anomaly_result: AnomalyResult): + visualizer = Visualizer(label_colors={"Anomaly": (18, 52, 86)}) + assert isinstance(visualizer.render(mock_image, anomaly_result), Image.Image) + + def test_detection_renders_with_rgb_tuple(self, mock_image: Image, detection_result: DetectionResult): + visualizer = Visualizer(label_colors={"car": (18, 52, 86)}) + assert isinstance(visualizer.render(mock_image, detection_result), Image.Image) + + def test_classification_renders_with_rgb_tuple( + self, + mock_image: Image, + classification_result: ClassificationResult, + ): + visualizer = Visualizer(label_colors={"cat": (18, 52, 86)}) + assert isinstance(visualizer.render(mock_image, classification_result), Image.Image) + + def test_tuple_and_hex_render_identically( + self, + mock_image: Image, + instance_segmentation_result: InstanceSegmentationResult, + ): + from_tuple = Visualizer(label_colors={"car": (18, 52, 86)}).render( + mock_image.copy(), + instance_segmentation_result, + ) + from_hex = Visualizer(label_colors={"car": "#123456"}).render( + mock_image.copy(), + instance_segmentation_result, + ) + assert from_tuple.tobytes() == from_hex.tobytes() + + +class TestToRgb: + """Tests for to_rgb().""" + + @pytest.mark.parametrize( + ("color", "expected"), + [ + ("#123456", (18, 52, 86)), + ("red", (255, 0, 0)), + ((18, 52, 86), (18, 52, 86)), + ], + ) + def test_normalizes_colors(self, color, expected): + assert to_rgb(color) == expected + + +class TestRenderedOutput: + """Regression tests on the rendered images.""" + + def test_empty_mapping_renders_like_no_mapping(self, mock_image: Image, detection_result: DetectionResult): + default = Visualizer().render(mock_image.copy(), detection_result) + empty = Visualizer(label_colors={}).render(mock_image.copy(), detection_result) + assert default.tobytes() == empty.tobytes() + + def test_mapping_changes_rendered_pixels(self, mock_image: Image, detection_result: DetectionResult): + default = Visualizer().render(mock_image.copy(), detection_result) + custom = Visualizer(label_colors={"car": "#123456"}).render(mock_image.copy(), detection_result) + assert default.tobytes() != custom.tobytes()