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
25 changes: 25 additions & 0 deletions src/model_api/visualizer/defaults.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
"""Default visualization constants."""

# Copyright (C) 2026 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

# Font sizes
DEFAULT_FONT_SIZE: int = 16
"""Default font size used for all text (labels, bounding boxes, overlays, keypoints)."""

# Line / outline widths
DEFAULT_OUTLINE_WIDTH: int = 2
"""Default outline width for bounding boxes and polygon contours."""

# Opacity
DEFAULT_OPACITY: float = 0.4
"""Default blend opacity for overlays and polygon fills."""

# Keypoint drawing
DEFAULT_KEYPOINT_SIZE: int = 3
"""Default radius (in pixels) for keypoint dots."""

# Scale baseline
SCALE_BASELINE: int = 1280
"""Longer-edge pixel count of 720p (landscape). Used as the denominator when
computing the auto-scale factor."""
6 changes: 5 additions & 1 deletion src/model_api/visualizer/layout/hstack.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,11 @@ def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, sc
for _primitive in scene.get_primitives(primitive):
image_ = _primitive.compute(image.copy())
if isinstance(_primitive, Overlay):
image_ = Overlay.overlay_labels(image=image_, labels=_primitive.label)
image_ = Overlay.overlay_labels(
image=image_,
labels=_primitive.label,
font_size=_primitive.font_size,
)
images.append(image_)
return self._stitch(*images)
return None
Expand Down
19 changes: 14 additions & 5 deletions src/model_api/visualizer/primitive/bounding_box.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@

from __future__ import annotations

from PIL import Image, ImageDraw
from PIL import Image, ImageDraw, ImageFont

from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_OUTLINE_WIDTH

from .primitive import Primitive

Expand All @@ -20,6 +22,8 @@ class BoundingBox(Primitive):
y2 (int): y-coordinate of the bottom-right corner of the bounding box.
label (str | None): Label of the bounding box.
color (str | tuple[int, int, int]): Color of the bounding box.
outline_width (int): Width of the bounding box outline.
font_size (int): Font size for the label text.

Example:
>>> bounding_box = BoundingBox(x1=10, y1=10, x2=100, y2=100, label="Label Name")
Expand All @@ -34,30 +38,35 @@ def __init__(
y2: int,
label: str | None = None,
color: str | tuple[int, int, int] = "blue",
outline_width: int = DEFAULT_OUTLINE_WIDTH,
font_size: int = DEFAULT_FONT_SIZE,
) -> None:
self.x1 = x1
self.y1 = y1
self.x2 = x2
self.y2 = y2
self.label = label
self.color = color
self.y_buffer = 5 # Text at the bottom of the text box is clipped. This prevents that.
self.outline_width = outline_width
self.font_size = font_size
self.font = ImageFont.load_default(size=self.font_size)
self.y_buffer = max(3, font_size // 3) # Text at the bottom of the text box is clipped. This prevents that.

def compute(self, image: Image) -> Image:
draw = ImageDraw.Draw(image)
# draw rectangle
draw.rectangle((self.x1, self.y1, self.x2, self.y2), outline=self.color, width=2)
draw.rectangle((self.x1, self.y1, self.x2, self.y2), outline=self.color, width=self.outline_width)
# add label
if self.label:
# draw the background of the label
textbox = draw.textbbox((0, 0), self.label)
textbox = draw.textbbox((0, 0), self.label, font=self.font)
label_image = Image.new(
"RGB",
(textbox[2] - textbox[0], textbox[3] + self.y_buffer - textbox[1]),
self.color,
)
draw = ImageDraw.Draw(label_image)
# write the label on the background
draw.text((0, 0), self.label, fill="white")
draw.text((0, 0), self.label, font=self.font, fill="white")
image.paste(label_image, (self.x1, self.y1))
return image
8 changes: 6 additions & 2 deletions src/model_api/visualizer/primitive/keypoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_KEYPOINT_SIZE

from .primitive import Primitive


Expand All @@ -25,12 +27,14 @@ def __init__(
keypoints: np.ndarray,
scores: Union[np.ndarray, None] = None,
color: Union[str, tuple[int, int, int]] = "purple",
keypoint_size: int = 3,
keypoint_size: int = DEFAULT_KEYPOINT_SIZE,
font_size: int = DEFAULT_FONT_SIZE,
) -> None:
self.keypoints = self._validate_keypoints(keypoints)
self.scores = scores
self.color = color
self.keypoint_size = keypoint_size
self.font_size = font_size

def compute(self, image: Image) -> Image:
"""Draw keypoints on the image."""
Expand All @@ -47,7 +51,7 @@ def compute(self, image: Image) -> Image:
)

if self.scores is not None:
font = ImageFont.load_default(size=18)
font = ImageFont.load_default(size=self.font_size)
for score, keypoint in zip(self.scores, self.keypoints):
textbox = draw.textbbox((0, 0), f"{score:.2f}", font=font)
draw.text(
Expand Down
4 changes: 3 additions & 1 deletion src/model_api/visualizer/primitive/label.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

from PIL import Image, ImageDraw, ImageFont

from model_api.visualizer.defaults import DEFAULT_FONT_SIZE

from .primitive import Primitive


Expand Down Expand Up @@ -46,7 +48,7 @@ def __init__(
fg_color: Union[str, tuple[int, int, int]] = "black",
bg_color: Union[str, tuple[int, int, int]] = "yellow",
font_path: Union[str, BytesIO, None] = None,
size: int = 16,
size: int = DEFAULT_FONT_SIZE,
) -> None:
self.label = f"{label} ({score:.2f})" if score is not None else label
self.fg_color = fg_color
Expand Down
22 changes: 18 additions & 4 deletions src/model_api/visualizer/primitive/overlay.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
import PIL
from PIL import ImageFont

from model_api.visualizer.defaults import DEFAULT_FONT_SIZE, DEFAULT_OPACITY

from .primitive import Primitive


Expand All @@ -28,12 +30,14 @@ class Overlay(Primitive):
def __init__(
self,
image: PIL.Image | np.ndarray,
opacity: float = 0.4,
opacity: float = DEFAULT_OPACITY,
label: Union[str, None] = None,
font_size: int = DEFAULT_FONT_SIZE,
) -> None:
self.image = self._to_pil(image)
self.label = label
self.opacity = opacity
self.font_size = font_size

def _to_pil(self, image: PIL.Image | np.ndarray) -> PIL.Image:
if isinstance(image, np.ndarray):
Expand All @@ -45,15 +49,25 @@ def compute(self, image: PIL.Image) -> PIL.Image:
return PIL.Image.blend(image, image_, self.opacity)

@classmethod
def overlay_labels(cls, image: PIL.Image, labels: Union[list[str], str, None] = None) -> PIL.Image:
def overlay_labels(
cls,
image: PIL.Image,
labels: Union[list[str], str, None] = None,
font_size: int = DEFAULT_FONT_SIZE,
) -> PIL.Image:
"""Draw labels at the bottom center of the image.

This is handy when you want to add a label to the image.

Args:
image: Image to overlay the labels on.
labels: Labels to overlay.
font_size: Font size for the label text.
"""
if labels is not None:
labels = [labels] if isinstance(labels, str) else labels
font = ImageFont.load_default(size=18)
buffer_y = 5
font = ImageFont.load_default(size=font_size)
buffer_y = max(3, font_size // 3)
dummy_image = PIL.Image.new("RGB", (1, 1))
draw = PIL.ImageDraw.Draw(dummy_image)
textbox = draw.textbbox((0, 0), ", ".join(labels), font=font)
Expand Down
6 changes: 4 additions & 2 deletions src/model_api/visualizer/primitive/polygon.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@
import cv2
from PIL import Image, ImageColor, ImageDraw

from model_api.visualizer.defaults import DEFAULT_OPACITY, DEFAULT_OUTLINE_WIDTH

from .primitive import Primitive

if TYPE_CHECKING:
Expand Down Expand Up @@ -41,8 +43,8 @@ def __init__(
points: list[tuple[int, int]] | None = None,
mask: np.ndarray | None = None,
color: str | tuple[int, int, int] = "blue",
opacity: float = 0.4,
outline_width: int = 2,
opacity: float = DEFAULT_OPACITY,
outline_width: int = DEFAULT_OUTLINE_WIDTH,
) -> None:
self.points = self._get_points(points, mask)
self.color = color
Expand Down
40 changes: 34 additions & 6 deletions src/model_api/visualizer/scene/anomaly.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@
# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from itertools import starmap
from typing import 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.layout import Flatten, Layout
from model_api.visualizer.primitive import BoundingBox, Label, Overlay, Polygon

Expand All @@ -19,7 +19,14 @@
class AnomalyScene(Scene):
"""Anomaly Scene."""

def __init__(self, image: Image, result: AnomalyResult, layout: Union[Layout, None] = None) -> None:
def __init__(
self,
image: Image,
result: AnomalyResult,
layout: Union[Layout, None] = None,
scale: float = 1.0,
) -> None:
self.scale = scale
super().__init__(
base=image,
overlay=self._get_overlays(result),
Expand All @@ -32,23 +39,44 @@ def __init__(self, image: Image, result: AnomalyResult, layout: Union[Layout, No
def _get_overlays(self, result: AnomalyResult) -> list[Overlay]:
if result.anomaly_map is not None:
anomaly_map = cv2.cvtColor(result.anomaly_map, cv2.COLOR_BGR2RGB)
return [Overlay(anomaly_map)]
return [Overlay(anomaly_map, font_size=int(DEFAULT_FONT_SIZE * self.scale))]
return []

def _get_bounding_boxes(self, result: AnomalyResult) -> list[BoundingBox]:
if result.pred_boxes is not None:
return list(starmap(BoundingBox, result.pred_boxes))
return [
BoundingBox(
x1=box[0],
y1=box[1],
x2=box[2],
y2=box[3],
outline_width=max(1, int(DEFAULT_OUTLINE_WIDTH * self.scale)),
font_size=int(DEFAULT_FONT_SIZE * self.scale),
)
for box in result.pred_boxes
]
return []

def _get_labels(self, result: AnomalyResult) -> list[Label]:
labels = []
if result.pred_label is not None and result.pred_score is not None:
labels.append(Label(label=result.pred_label, score=result.pred_score))
labels.append(
Label(
label=result.pred_label,
score=result.pred_score,
size=int(DEFAULT_FONT_SIZE * self.scale),
),
)
return labels

def _get_polygons(self, result: AnomalyResult) -> list[Polygon]:
if result.pred_mask is not None:
return [Polygon(result.pred_mask)]
return [
Polygon(
result.pred_mask,
outline_width=max(1, int(DEFAULT_OUTLINE_WIDTH * self.scale)),
),
]
return []

@property
Expand Down
18 changes: 16 additions & 2 deletions src/model_api/visualizer/scene/classification.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from PIL import Image

from model_api.models.result import ClassificationResult
from model_api.visualizer.defaults import DEFAULT_FONT_SIZE
from model_api.visualizer.layout import Flatten, Layout
from model_api.visualizer.primitive import Label, Overlay

Expand All @@ -18,7 +19,14 @@
class ClassificationScene(Scene):
"""Classification Scene."""

def __init__(self, image: Image, result: ClassificationResult, layout: Union[Layout, None] = None) -> None:
def __init__(
self,
image: Image,
result: ClassificationResult,
layout: Union[Layout, None] = None,
scale: float = 1.0,
) -> None:
self.scale = scale
super().__init__(
base=image,
label=self._get_labels(result),
Expand All @@ -31,7 +39,13 @@ def _get_labels(self, result: ClassificationResult) -> list[Label]:
if result.top_labels is not None and len(result.top_labels) > 0:
for label in result.top_labels:
if label.name is not None:
labels.append(Label(label=label.name, score=label.confidence))
labels.append(
Label(
label=label.name,
score=label.confidence,
size=int(DEFAULT_FONT_SIZE * self.scale),
),
)
return labels

def _get_overlays(self, result: ClassificationResult) -> list[Overlay]:
Expand Down
Loading