Add Anomaly scene

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Author: ashwinvaidya17

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

src/python/model_api/visualizer/primitive/label.py

@@ -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)

src/python/model_api/visualizer/scene/anomaly.py

@@ -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)

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
 

src/python/model_api/visualizer/scene/detection.py

@@ -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

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)

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)

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()