diff --git a/src/python/model_api/visualizer/layout/flatten.py b/src/python/model_api/visualizer/layout/flatten.py index ba128658..bfc2bdf5 100644 --- a/src/python/model_api/visualizer/layout/flatten.py +++ b/src/python/model_api/visualizer/layout/flatten.py @@ -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 diff --git a/src/python/model_api/visualizer/primitive/label.py b/src/python/model_api/visualizer/primitive/label.py index 89dbfe09..8c77ec42 100644 --- a/src/python/model_api/visualizer/primitive/label.py +++ b/src/python/model_api/visualizer/primitive/label.py @@ -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. @@ -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) diff --git a/src/python/model_api/visualizer/scene/anomaly.py b/src/python/model_api/visualizer/scene/anomaly.py index 7bb02756..8790dbd5 100644 --- a/src/python/model_api/visualizer/scene/anomaly.py +++ b/src/python/model_api/visualizer/scene/anomaly.py @@ -3,11 +3,15 @@ # 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 @@ -15,10 +19,38 @@ 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) diff --git a/src/python/model_api/visualizer/scene/classification.py b/src/python/model_api/visualizer/scene/classification.py index 54f51ba7..ec2f3fe6 100644 --- a/src/python/model_api/visualizer/scene/classification.py +++ b/src/python/model_api/visualizer/scene/classification.py @@ -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 @@ -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 diff --git a/src/python/model_api/visualizer/scene/detection.py b/src/python/model_api/visualizer/scene/detection.py index 35e55bc2..2cc7691a 100644 --- a/src/python/model_api/visualizer/scene/detection.py +++ b/src/python/model_api/visualizer/scene/detection.py @@ -3,9 +3,12 @@ # 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 @@ -13,6 +16,6 @@ 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 diff --git a/src/python/model_api/visualizer/scene/scene.py b/src/python/model_api/visualizer/scene/scene.py index 161c41fc..d798a0da 100644 --- a/src/python/model_api/visualizer/scene/scene.py +++ b/src/python/model_api/visualizer/scene/scene.py @@ -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) diff --git a/src/python/model_api/visualizer/visualizer.py b/src/python/model_api/visualizer/visualizer.py index 8a2ea87e..e665e59a 100644 --- a/src/python/model_api/visualizer/visualizer.py +++ b/src/python/model_api/visualizer/visualizer.py @@ -4,6 +4,7 @@ # SPDX-License-Identifier: Apache-2.0 from pathlib import Path +from typing import Union from PIL import Image @@ -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: @@ -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) diff --git a/tests/python/unit/visualizer/test_scene.py b/tests/python/unit/visualizer/test_scene.py new file mode 100644 index 00000000..b25538e5 --- /dev/null +++ b/tests/python/unit/visualizer/test_scene.py @@ -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()