Add Overlay and Flatten

ashwinvaidya17 · ashwinvaidya17 · commit 42344e56bdc7 · 2024-12-20T17:52:58.000+01:00
Signed-off-by: Ashwin Vaidya &lt;ashwinnitinvaidya@gmail.com&gt;
diff --git a/model_api/python/model_api/models/__init__.py b/model_api/python/model_api/models/__init__.py
@@ -56,44 +56,47 @@
 
 __all__ = [
     "ActionClassificationModel",
+    "add_rotated_rects",
     "AnomalyDetection",
     "AnomalyResult",
+    "classification_models",
     "ClassificationModel",
+    "ClassificationResult",
     "Contour",
+    "detection_models",
+    "DetectedKeypoints",
     "DetectionModel",
+    "DetectionResult",
+    "get_contours",
     "ImageModel",
     "ImageResultWithSoftPrediction",
     "InstanceSegmentationResult",
-    "VisualPromptingResult",
-    "ZSLVisualPromptingResult",
-    "PredictedMask",
-    "SAMVisualPrompter",
-    "SAMLearnableVisualPrompter",
     "KeypointDetectionModel",
-    "TopDownKeypointDetectionPipeline",
+    "Label",
     "MaskRCNNModel",
     "Model",
     "OutputTransform",
+    "PredictedMask",
+    "Prompt",
+    "RotatedSegmentationResult",
+    "SAMDecoder",
+    "SAMImageEncoder",
+    "SAMLearnableVisualPrompter",
+    "SAMVisualPrompter",
     "SalientObjectDetectionModel",
+    "segmentation_models",
     "SegmentationModel",
     "SSD",
+    "TopDownKeypointDetectionPipeline",
+    "VisualPromptingResult",
     "YOLO",
-    "YoloV3ONNX",
-    "YoloV4",
+    "YOLOF",
+    "YOLOv3ONNX",
+    "YOLOv4",
     "YOLOv5",
     "YOLOv8",
-    "YOLOF",
     "YOLOX",
-    "SAMDecoder",
-    "SAMImageEncoder",
-    "ClassificationResult",
-    "Prompt",
-    "DetectionResult",
-    "DetectedKeypoints",
-    "classification_models",
-    "detection_models",
-    "segmentation_models",
-    "RotatedSegmentationResult",
-    "add_rotated_rects",
-    "get_contours",
+    "ZSLVisualPromptingResult",
+    "YoloV3ONNX",
+    "YoloV4",
 ]
diff --git a/model_api/python/model_api/models/action_classification.py b/model_api/python/model_api/models/action_classification.py
@@ -10,10 +10,9 @@
 import numpy as np
 
 from model_api.adapters.utils import RESIZE_TYPES, InputTransform
-from model_api.models.result import Label
+from model_api.models.result import ClassificationResult, Label
 
 from .model import Model
-from .result import ClassificationResult
 from .types import BooleanValue, ListValue, NumericalValue, StringValue
 from .utils import load_labels
 
diff --git a/model_api/python/model_api/models/result/__init__.py b/model_api/python/model_api/models/result/__init__.py
@@ -3,6 +3,14 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from .scene import (
+    AnomalyScene,
+    ClassificationScene,
+    DetectionScene,
+    KeypointScene,
+    SegmentationScene,
+    VisualPromptingScene,
+)
 from .types import (
     AnomalyResult,
     ClassificationResult,
@@ -33,4 +41,10 @@
     "VisualPromptingResult",
     "ZSLVisualPromptingResult",
     "RotatedSegmentationResult",
+    "AnomalyScene",
+    "ClassificationScene",
+    "DetectionScene",
+    "KeypointScene",
+    "SegmentationScene",
+    "VisualPromptingScene",
 ]
diff --git a/model_api/python/model_api/models/result/scene/__init__.py b/model_api/python/model_api/models/result/scene/__init__.py
@@ -2,3 +2,19 @@
 
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
+
+from .anomaly import AnomalyScene
+from .classification import ClassificationScene
+from .detection import DetectionScene
+from .keypoint import KeypointScene
+from .segmentation import SegmentationScene
+from .visual_prompting import VisualPromptingScene
+
+__all__ = [
+    "AnomalyScene",
+    "ClassificationScene",
+    "DetectionScene",
+    "KeypointScene",
+    "SegmentationScene",
+    "VisualPromptingScene",
+]
diff --git a/model_api/python/model_api/models/result/scene/anomaly.py b/model_api/python/model_api/models/result/scene/anomaly.py
@@ -0,0 +1,23 @@
+"""Anomaly Scene."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from PIL import Image
+
+from model_api.models.result.types import AnomalyResult
+from model_api.visualizer.layout import Flatten, Layout
+from model_api.visualizer.primitive import Overlay
+from model_api.visualizer.scene import Scene
+
+
+class AnomalyScene(Scene):
+    """Anomaly Scene."""
+
+    def __init__(self, image: Image, result: AnomalyResult) -> None:
+        self.image = image
+        self.result = result
+
+    @property
+    def default_layout(self) -> Layout:
+        return Flatten(Overlay)
diff --git a/model_api/python/model_api/models/result/scene/classification.py b/model_api/python/model_api/models/result/scene/classification.py
@@ -0,0 +1,23 @@
+"""Classification Scene."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from PIL import Image
+
+from model_api.models.result.types import ClassificationResult
+from model_api.visualizer.layout import Flatten, Layout
+from model_api.visualizer.primitive import Overlay
+from model_api.visualizer.scene import Scene
+
+
+class ClassificationScene(Scene):
+    """Classification Scene."""
+
+    def __init__(self, image: Image, result: ClassificationResult) -> None:
+        self.image = image
+        self.result = result
+
+    @property
+    def default_layout(self) -> Layout:
+        return Flatten(Overlay)
diff --git a/model_api/python/model_api/models/result/scene/detection.py b/model_api/python/model_api/models/result/scene/detection.py
@@ -0,0 +1,17 @@
+"""Detection Scene."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from PIL import Image
+
+from model_api.models.result.types import DetectionResult
+from model_api.visualizer import Scene
+
+
+class DetectionScene(Scene):
+    """Detection Scene."""
+
+    def __init__(self, image: Image, result: DetectionResult) -> None:
+        self.image = image
+        self.result = result
diff --git a/model_api/python/model_api/models/result/scene/keypoint.py b/model_api/python/model_api/models/result/scene/keypoint.py
@@ -0,0 +1,20 @@
+"""Keypoint Scene."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from model_api.models.result.types import DetectedKeypoints
+from model_api.visualizer.layout import Flatten, Layout
+from model_api.visualizer.primitive import Overlay
+from model_api.visualizer.scene import Scene
+
+
+class KeypointScene(Scene):
+    """Keypoint Scene."""
+
+    def __init__(self, result: DetectedKeypoints) -> None:
+        self.result = result
+
+    @property
+    def default_layout(self) -> Layout:
+        return Flatten(Overlay)
diff --git a/model_api/python/model_api/models/result/scene/segmentation.py b/model_api/python/model_api/models/result/scene/segmentation.py
@@ -0,0 +1,14 @@
+"""Segmentation Scene."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from model_api.models.result.types import InstanceSegmentationResult
+from model_api.visualizer.scene import Scene
+
+
+class SegmentationScene(Scene):
+    """Segmentation Scene."""
+
+    def __init__(self, result: InstanceSegmentationResult) -> None:
+        self.result = result
diff --git a/model_api/python/model_api/models/result/scene/visual_prompting.py b/model_api/python/model_api/models/result/scene/visual_prompting.py
@@ -0,0 +1,14 @@
+"""Visual Prompting Scene."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from model_api.models.result.types import VisualPromptingResult
+from model_api.visualizer import Scene
+
+
+class VisualPromptingScene(Scene):
+    """Visual Prompting Scene."""
+
+    def __init__(self, result: VisualPromptingResult) -> None:
+        self.result = result
diff --git a/model_api/python/model_api/visualizer/__init__.py b/model_api/python/model_api/visualizer/__init__.py
@@ -3,7 +3,9 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from .layout import Flatten, Layout
+from .primitive import Overlay
 from .scene import Scene
 from .visualizer import Visualizer
 
-__all__ = ["Scene", "Visualizer"]
+__all__ = ["Overlay", "Scene", "Visualizer", "Layout", "Flatten"]
diff --git a/model_api/python/model_api/visualizer/layout.py b/model_api/python/model_api/visualizer/layout.py
@@ -5,8 +5,8 @@
 
 from __future__ import annotations
 
-from abc import ABC
-from typing import TYPE_CHECKING
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Type
 
 if TYPE_CHECKING:
     import PIL
@@ -19,10 +19,31 @@
 class Layout(ABC):
     """Base class for layouts."""
 
-    def _compute_on_primitive(self, primitive: Primitive, image: PIL.Image, scene: Scene) -> PIL.Image | None:
-        if scene.has_primitives(type(primitive)):
-            primitives = scene.get_primitives(type(primitive))
-            for primitive in primitives:
-                image = primitive.compute(image)
+    def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, scene: Scene) -> PIL.Image | None:
+        if scene.has_primitives(primitive):
+            primitives = scene.get_primitives(primitive)
+            for _primitive in primitives:
+                image = _primitive.compute(image)
             return image
         return None
+
+    @abstractmethod
+    def __call__(self, scene: Scene) -> PIL.Image:
+        """Compute the layout."""
+
+
+class Flatten(Layout):
+    """Put all primitives on top of each other.
+
+    Args:
+        *args (Type[Primitive]): Primitives to be applied.
+    """
+
+    def __init__(self, *args: Type[Primitive]) -> None:
+        self.children = args
+
+    def __call__(self, scene: Scene) -> PIL.Image:
+        _image: PIL.Image = scene.base.copy()
+        for child in self.children:
+            _image = self._compute_on_primitive(child, _image, scene)
+        return _image
diff --git a/model_api/python/model_api/visualizer/primitive.py b/model_api/python/model_api/visualizer/primitive.py
@@ -6,15 +6,38 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
 
-if TYPE_CHECKING:
-    import PIL
+import numpy as np
+import PIL
 
 
 class Primitive(ABC):
     """Primitive class."""
 
     @abstractmethod
-    def compute(self, image: PIL.Image, **kwargs) -> PIL.Image:
+    def compute(self, image: PIL.Image) -> PIL.Image:
         pass
+
+
+class Overlay(Primitive):
+    """Overlay primitive.
+
+    Useful for XAI and Anomaly Maps.
+
+    Args:
+        image (PIL.Image | np.ndarray): Image to be overlaid.
+        opacity (float): Opacity of the overlay.
+    """
+
+    def __init__(self, image: PIL.Image | np.ndarray, opacity: float = 0.4) -> None:
+        self.image = self._to_pil(image)
+        self.opacity = opacity
+
+    def _to_pil(self, image: PIL.Image | np.ndarray) -> PIL.Image:
+        if isinstance(image, np.ndarray):
+            return PIL.Image.fromarray(image)
+        return image
+
+    def compute(self, image: PIL.Image) -> PIL.Image:
+        _image = self.image.resize(image.size)
+        return PIL.Image.blend(image, _image, self.opacity)
diff --git a/model_api/python/model_api/visualizer/scene.py b/model_api/python/model_api/visualizer/scene.py
diff --git a/model_api/python/model_api/visualizer/visualizer.py b/model_api/python/model_api/visualizer/visualizer.py
diff --git a/tests/python/unit/visualizer/__init__.py b/tests/python/unit/visualizer/__init__.py
diff --git a/tests/python/unit/visualizer/conftest.py b/tests/python/unit/visualizer/conftest.py
diff --git a/tests/python/unit/visualizer/test_layout.py b/tests/python/unit/visualizer/test_layout.py
diff --git a/tests/python/unit/visualizer/test_primitive.py b/tests/python/unit/visualizer/test_primitive.py
