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
8 changes: 5 additions & 3 deletions src/python/model_api/visualizer/layout/flatten.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,9 @@ def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, sc
return None

def __call__(self, scene: Scene) -> PIL.Image:
image_: PIL.Image = scene.base.copy()
image: PIL.Image = scene.base.copy()
for child in self.children:
image_ = child(scene) if isinstance(child, Layout) else self._compute_on_primitive(child, image_, scene)
return image_
image_ = child(scene) if isinstance(child, Layout) else self._compute_on_primitive(child, image, scene)
if image_ is not None:
image = image_
return image
4 changes: 3 additions & 1 deletion src/python/model_api/visualizer/primitive/label.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ class Label(Primitive):

Args:
label (str): Text of the label.
score (float | None): Score of the label. This is optional.
fg_color (str | tuple[int, int, int]): Foreground color of the label.
bg_color (str | tuple[int, int, int]): Background color of the label.
font_path (str | None | BytesIO): Path to the font file.
Expand All @@ -41,12 +42,13 @@ class Label(Primitive):
def __init__(
self,
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",
font_path: Union[str, BytesIO, None] = None,
size: int = 16,
) -> None:
self.label = label
self.label = f"{label} ({score:.2f})" if score is not None else label
self.fg_color = fg_color
self.bg_color = bg_color
self.font = ImageFont.load_default(size=size) if font_path is None else ImageFont.truetype(font_path, size)
Expand Down
42 changes: 37 additions & 5 deletions src/python/model_api/visualizer/scene/anomaly.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,22 +3,54 @@
# 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.layout import Flatten, Layout
from model_api.visualizer.primitive import Overlay
from model_api.visualizer.primitive import BoundingBox, Label, Overlay, Polygon

from .scene import Scene


class AnomalyScene(Scene):
"""Anomaly Scene."""

def __init__(self, image: Image, result: AnomalyResult) -> None:
self.image = image
self.result = result
def __init__(self, image: Image, result: AnomalyResult, layout: Union[Layout, None] = None) -> None:
super().__init__(
base=image,
overlay=self._get_overlays(result),
bounding_box=self._get_bounding_boxes(result),
label=self._get_labels(result),
polygon=self._get_polygons(result),
layout=layout,
)

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 []

def _get_bounding_boxes(self, result: AnomalyResult) -> list[BoundingBox]:
if result.pred_boxes is not None:
return list(starmap(BoundingBox, 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))
return labels

def _get_polygons(self, result: AnomalyResult) -> list[Polygon]:
if result.pred_mask is not None:
return [Polygon(result.pred_mask)]
return []

@property
def default_layout(self) -> Layout:
return Flatten(Overlay)
return Flatten(Overlay, BoundingBox, Label, Polygon)
4 changes: 3 additions & 1 deletion src/python/model_api/visualizer/scene/classification.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from typing import Union

from PIL import Image

from model_api.models.result import ClassificationResult
Expand All @@ -15,7 +17,7 @@
class ClassificationScene(Scene):
"""Classification Scene."""

def __init__(self, image: Image, result: ClassificationResult) -> None:
def __init__(self, image: Image, result: ClassificationResult, layout: Union[Layout, None] = None) -> None:
self.image = image
self.result = result

Expand Down
5 changes: 4 additions & 1 deletion src/python/model_api/visualizer/scene/detection.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,19 @@
# Copyright (C) 2024 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

from typing import Union

from PIL import Image

from model_api.models.result import DetectionResult
from model_api.visualizer.layout import Layout

from .scene import Scene


class DetectionScene(Scene):
"""Detection Scene."""

def __init__(self, image: Image, result: DetectionResult) -> None:
def __init__(self, image: Image, result: DetectionResult, layout: Union[Layout, None] = None) -> None:
self.image = image
self.result = result
3 changes: 2 additions & 1 deletion src/python/model_api/visualizer/scene/scene.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,8 @@ def __init__(
self.polygon = self._to_polygon(polygon)
self.layout = layout

def show(self) -> Image: ...
def show(self) -> None:
self.render().show()

def save(self, path: Path) -> None:
self.render().save(path)
Expand Down
11 changes: 7 additions & 4 deletions src/python/model_api/visualizer/visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
# SPDX-License-Identifier: Apache-2.0

from pathlib import Path
from typing import Union

from PIL import Image

Expand All @@ -19,7 +20,9 @@


class Visualizer:
def __init__(self, layout: Layout) -> None:
"""Utility class to automatically select the correct scene and render/show it."""

def __init__(self, layout: Union[Layout, None] = None) -> None:
self.layout = layout

def show(self, image: Image, result: Result) -> Image:
Expand All @@ -33,11 +36,11 @@ def save(self, image: Image, result: Result, path: Path) -> None:
def _scene_from_result(self, image: Image, result: Result) -> Scene:
scene: Scene
if isinstance(result, AnomalyResult):
scene = AnomalyScene(image, result)
scene = AnomalyScene(image, result, self.layout)
elif isinstance(result, ClassificationResult):
scene = ClassificationScene(image, result)
scene = ClassificationScene(image, result, self.layout)
elif isinstance(result, DetectionResult):
scene = DetectionScene(image, result)
scene = DetectionScene(image, result, self.layout)
else:
msg = f"Unsupported result type: {type(result)}"
raise ValueError(msg)
Expand Down
34 changes: 34 additions & 0 deletions tests/python/unit/visualizer/test_scene.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
"""Tests for scene."""

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

from pathlib import Path

import numpy as np
from PIL import Image

from model_api.models.result import AnomalyResult
from model_api.visualizer import Visualizer


def test_anomaly_scene(mock_image: Image, tmpdir: Path):
"""Test if the anomaly scene is created."""
heatmap = np.ones(mock_image.size, dtype=np.uint8)
heatmap *= 255

mask = np.zeros(mock_image.size, dtype=np.uint8)
mask[32:96, 32:96] = 255
mask[40:80, 0:128] = 255

anomaly_result = AnomalyResult(
anomaly_map=heatmap,
pred_boxes=np.array([[0, 0, 128, 128], [32, 32, 96, 96]]),
pred_label="Anomaly",
pred_mask=mask,
pred_score=0.85,
)

visualizer = Visualizer()
visualizer.save(mock_image, anomaly_result, tmpdir / "anomaly_scene.jpg")
assert Path(tmpdir / "anomaly_scene.jpg").exists()
Loading