diff --git a/src/model_api/visualizer/primitive/polygon.py b/src/model_api/visualizer/primitive/polygon.py index 74357264..2b30e21f 100644 --- a/src/model_api/visualizer/primitive/polygon.py +++ b/src/model_api/visualizer/primitive/polygon.py @@ -69,8 +69,8 @@ def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | N elif mask is not None: points_ = self._get_points_from_mask(mask) else: - msg = "Either points or mask should be provided." - raise ValueError(msg) + logger.warning("Neither points nor mask provided. Skipping polygon drawing.") + return [] return points_ def _get_points_from_mask(self, mask: np.ndarray) -> list[tuple[int, int]]: @@ -88,8 +88,8 @@ def _get_points_from_mask(self, mask: np.ndarray) -> list[tuple[int, int]]: logger.warning("Multiple contours found in the mask. Using the largest one.") contours = sorted(contours, key=cv2.contourArea, reverse=True) if len(contours) == 0: - msg = "No contours found in the mask." - raise ValueError(msg) + logger.warning("No contours found in the mask. Skipping polygon drawing.") + return [] points_ = contours[0].squeeze().tolist() return [tuple(point) for point in points_] @@ -102,6 +102,9 @@ def compute(self, image: Image) -> Image: Returns: Image with the polygon drawn on it. """ + if len(self.points) == 0: + return image + draw = ImageDraw.Draw(image, "RGBA") # Draw polygon with darker edge and a semi-transparent fill. ink = ImageColor.getrgb(self.color) diff --git a/src/model_api/visualizer/scene/detection.py b/src/model_api/visualizer/scene/detection.py index 7ffa3e97..5bf31847 100644 --- a/src/model_api/visualizer/scene/detection.py +++ b/src/model_api/visualizer/scene/detection.py @@ -13,12 +13,14 @@ from model_api.visualizer.primitive import BoundingBox, Label, Overlay from .scene import Scene +from .utils import get_label_color_mapping class DetectionScene(Scene): """Detection Scene.""" def __init__(self, image: Image, result: DetectionResult, layout: Union[Layout, None] = None) -> None: + self.color_per_label = get_label_color_mapping(result.label_names) super().__init__( base=image, bounding_box=self._get_bounding_boxes(result), @@ -43,7 +45,9 @@ def _get_bounding_boxes(self, result: DetectionResult) -> list[BoundingBox]: for score, label_name, bbox in zip(result.scores, result.label_names, result.bboxes): x1, y1, x2, y2 = bbox label = f"{label_name} ({score:.2f})" - bounding_boxes.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, label=label)) + bounding_boxes.append( + BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, label=label, color=self.color_per_label[label_name]), + ) return bounding_boxes @property diff --git a/src/model_api/visualizer/scene/segmentation/instance_segmentation.py b/src/model_api/visualizer/scene/segmentation/instance_segmentation.py index 8c89ca01..df3bf8f9 100644 --- a/src/model_api/visualizer/scene/segmentation/instance_segmentation.py +++ b/src/model_api/visualizer/scene/segmentation/instance_segmentation.py @@ -3,7 +3,6 @@ # Copyright (C) 2025 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -import random from typing import Union import cv2 @@ -13,15 +12,14 @@ from model_api.visualizer.layout import Flatten, HStack, Layout from model_api.visualizer.primitive import BoundingBox, Label, Overlay, Polygon from model_api.visualizer.scene import Scene +from model_api.visualizer.scene.utils import get_label_color_mapping class InstanceSegmentationScene(Scene): """Instance Segmentation Scene.""" def __init__(self, image: Image, result: InstanceSegmentationResult, layout: Union[Layout, None] = None) -> None: - # nosec as random is used for color generation - g = random.Random(0) # noqa: S311 # nosec B311 - self.color_per_label = {label: f"#{g.randint(0, 0xFFFFFF):06x}" for label in set(result.label_names)} # nosec B311 + self.color_per_label = get_label_color_mapping(result.label_names) super().__init__( base=image, label=self._get_labels(result), diff --git a/src/model_api/visualizer/scene/utils.py b/src/model_api/visualizer/scene/utils.py new file mode 100644 index 00000000..9ad97723 --- /dev/null +++ b/src/model_api/visualizer/scene/utils.py @@ -0,0 +1,40 @@ +"""Visualizer utilities.""" + +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +COLOR_PALETTE = [ + "#FF6B6B", # Red + "#4ECDC4", # Teal + "#45B7D1", # Blue + "#FFA07A", # Light Salmon + "#98D8C8", # Mint + "#F7DC6F", # Yellow + "#BB8FCE", # Purple + "#85C1E2", # Sky Blue + "#F8B739", # Orange + "#52BE80", # Green + "#EC7063", # Coral + "#5DADE2", # Light Blue + "#F39C12", # Dark Orange + "#8E44AD", # Dark Purple + "#16A085", # Dark Teal + "#E74C3C", # Dark Red + "#3498DB", # Dodger Blue + "#2ECC71", # Emerald + "#F1C40F", # Sun Yellow + "#E67E22", # Carrot Orange +] + + +def get_label_color_mapping(labels: list[str]) -> dict[str, str]: + """Generate a consistent color mapping for a list of labels. + + Args: + labels: List of label names. + + Returns: + Dictionary mapping each label to a hex color string. + """ + unique_labels = sorted(set(labels)) + return {label: COLOR_PALETTE[i % len(COLOR_PALETTE)] for i, label in enumerate(unique_labels)} diff --git a/tests/unit/visualizer/test_primitive.py b/tests/unit/visualizer/test_primitive.py index ef965658..55118cbd 100644 --- a/tests/unit/visualizer/test_primitive.py +++ b/tests/unit/visualizer/test_primitive.py @@ -5,7 +5,6 @@ import numpy as np import PIL -import pytest from PIL import ImageDraw from model_api.visualizer import BoundingBox, Keypoint, Label, Overlay, Polygon @@ -58,8 +57,11 @@ def test_polygon(mock_image: PIL.Image): polygon = Polygon(mask=mask, color="red", opacity=1, outline_width=1) assert polygon.compute(mock_image) == expected_image - with pytest.raises(ValueError, match="No contours found in the mask."): - Polygon(mask=np.zeros((100, 100), dtype=np.uint8)).compute(mock_image) + # Test with empty mask - should not raise, just return image unchanged + empty_mask = np.zeros((100, 100), dtype=np.uint8) + polygon_empty = Polygon(mask=empty_mask) + result = polygon_empty.compute(mock_image) + assert result == mock_image def test_label(mock_image: PIL.Image):