Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 7 additions & 4 deletions src/model_api/visualizer/primitive/polygon.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand All @@ -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_]

Expand All @@ -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)
Expand Down
6 changes: 5 additions & 1 deletion src/model_api/visualizer/scene/detection.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand All @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
# Copyright (C) 2025 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import random
from typing import Union

import cv2
Expand All @@ -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),
Expand Down
40 changes: 40 additions & 0 deletions src/model_api/visualizer/scene/utils.py
Original file line number Diff line number Diff line change
@@ -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)}
8 changes: 5 additions & 3 deletions tests/unit/visualizer/test_primitive.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand Down