Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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: 2 additions & 2 deletions src/python/model_api/visualizer/primitive/label.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,13 @@ class Label(Primitive):

def __init__(
self,
label: str,
label: Union[str, float],
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 = str(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
44 changes: 39 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,56 @@
# 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:
labels.append(Label(result.pred_label))
if result.pred_score is not None:
labels.append(Label(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