[WIP] Visualization POC 3

model_api/python/model_api/models/__init__.py

@@ -11,7 +11,7 @@
 from .instance_segmentation import MaskRCNNModel
 from .keypoint_detection import KeypointDetectionModel, TopDownKeypointDetectionPipeline
 from .model import Model
-from .result_types import (
+from .result import (
     AnomalyResult,
     ClassificationResult,
     Contour,

model_api/python/model_api/models/action_classification.py

@@ -12,7 +12,7 @@
 from model_api.adapters.utils import RESIZE_TYPES, InputTransform
 
 from .model import Model
-from .result_types import ClassificationResult
+from .result import ClassificationResult
 from .types import BooleanValue, ListValue, NumericalValue, StringValue
 from .utils import load_labels
 

model_api/python/model_api/models/anomaly.py

@@ -15,7 +15,7 @@
 import numpy as np
 
 from model_api.models.image_model import ImageModel
-from model_api.models.result_types import AnomalyResult
+from model_api.models.result import AnomalyResult
 from model_api.models.types import ListValue, NumericalValue, StringValue
 
 if TYPE_CHECKING:

model_api/python/model_api/models/classification.py

@@ -17,7 +17,7 @@
 from openvino.runtime import opset10 as opset
 
 from model_api.models.image_model import ImageModel
-from model_api.models.result_types import ClassificationResult
+from model_api.models.result import ClassificationResult
 from model_api.models.types import BooleanValue, ListValue, NumericalValue, StringValue
 from model_api.models.utils import softmax
 

model_api/python/model_api/models/detection_model.py

@@ -4,7 +4,7 @@
 #
 
 from .image_model import ImageModel
-from .result_types import Detection
+from .result import Detection
 from .types import ListValue, NumericalValue, StringValue
 from .utils import load_labels
 

model_api/python/model_api/models/instance_segmentation.py

@@ -9,7 +9,7 @@
 from model_api.adapters.inference_adapter import InferenceAdapter
 
 from .image_model import ImageModel
-from .result_types import InstanceSegmentationResult, SegmentedObject
+from .result import InstanceSegmentationResult, SegmentedObject
 from .types import BooleanValue, ListValue, NumericalValue, StringValue
 from .utils import load_labels
 

model_api/python/model_api/models/keypoint_detection.py

@@ -10,7 +10,7 @@
 import numpy as np
 
 from .image_model import ImageModel
-from .result_types import DetectedKeypoints, Detection
+from .result import DetectedKeypoints, Detection
 from .types import ListValue
 
 

model_api/python/model_api/models/result/__init__.py

@@ -0,0 +1,36 @@
+"""Model results."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .types import (
+    AnomalyResult,
+    ClassificationResult,
+    Contour,
+    DetectedKeypoints,
+    Detection,
+    DetectionResult,
+    ImageResultWithSoftPrediction,
+    InstanceSegmentationResult,
+    PredictedMask,
+    SegmentedObject,
+    SegmentedObjectWithRects,
+    VisualPromptingResult,
+    ZSLVisualPromptingResult,
+)
+
+__all__ = [
+    "AnomalyResult",
+    "ClassificationResult",
+    "Contour",
+    "Detection",
+    "DetectionResult",
+    "DetectedKeypoints",
+    "SegmentedObject",
+    "SegmentedObjectWithRects",
+    "ImageResultWithSoftPrediction",
+    "InstanceSegmentationResult",
+    "PredictedMask",
+    "VisualPromptingResult",
+    "ZSLVisualPromptingResult",
+]

model_api/python/model_api/models/result/media/__init__.py

@@ -0,0 +1,10 @@
+"""Result visualization media."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .anomaly import AnomalyMedia
+
+__all__ = [
+    "AnomalyMedia",
+]

model_api/python/model_api/models/result/media/anomaly.py

@@ -0,0 +1,23 @@
+"""Anomaly result media."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import cv2
+
+from model_api.models.result.types import AnomalyResult
+from model_api.visualizer.layout import Flatten, Layout
+from model_api.visualizer.media import Media
+from model_api.visualizer.primitives import Overlay
+
+
+class AnomalyMedia(Media):
+    """Anomaly result media."""
+
+    def __init__(self, result: AnomalyResult) -> None:
+        anomaly_map = cv2.applyColorMap(result.anomaly_map, cv2.COLORMAP_JET)
+        super().__init__(Overlay(anomaly_map))
+
+    @property
+    def default_layout(self) -> Layout:
+        return Flatten(Overlay)

model_api/python/model_api/models/result_types/__init__.py → model_api/python/model_api/models/result/types/__init__.py

@@ -1,4 +1,4 @@
-"""Result types."""
+"""Result containers."""
 
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0

model_api/python/model_api/models/result_types/anomaly.py → model_api/python/model_api/models/result/types/anomaly.py

@@ -5,10 +5,15 @@
 
 from __future__ import annotations
 
+import cv2
 import numpy as np
 
+from model_api.visualizer.layout import Flatten, Layout
 
-class AnomalyResult:
+from .base import Result
+
+
+class AnomalyResult(Result):
     """Results for anomaly models."""
 
     def __init__(
@@ -19,6 +24,7 @@ def __init__(
         pred_mask: np.ndarray | None = None,
         pred_score: float | None = None,
     ) -> None:
+        super().__init__()
         self.anomaly_map = anomaly_map
         self.pred_boxes = pred_boxes
         self.pred_label = pred_label

model_api/python/model_api/models/result/types/base.py

@@ -0,0 +1,10 @@
+"""Base result type"""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import ABC
+
+
+class Result(ABC):
+    """Base result type."""

model_api/python/model_api/models/result_types/classification.py → model_api/python/model_api/models/result/types/classification.py

@@ -7,13 +7,14 @@
 
 from typing import TYPE_CHECKING
 
+from .base import Result
 from .utils import array_shape_to_str
 
 if TYPE_CHECKING:
     import numpy as np
 
 
-class ClassificationResult:
+class ClassificationResult(Result):
     """Results for classification models."""
 
     def __init__(

model_api/python/model_api/models/segmentation.py

@@ -11,7 +11,7 @@
 import numpy as np
 
 from model_api.models.image_model import ImageModel
-from model_api.models.result_types import Contour, ImageResultWithSoftPrediction
+from model_api.models.result import Contour, ImageResultWithSoftPrediction
 from model_api.models.types import BooleanValue, ListValue, NumericalValue, StringValue
 from model_api.models.utils import load_labels
 

model_api/python/model_api/models/ssd.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from .detection_model import DetectionModel
-from .result_types import Detection, DetectionResult
+from .result import Detection, DetectionResult
 
 
 class SSD(DetectionModel):

model_api/python/model_api/models/utils.py

@@ -10,7 +10,7 @@
 import cv2
 import numpy as np
 
-from model_api.models.result_types import Contour, Detection, SegmentedObject, SegmentedObjectWithRects
+from model_api.models.result import Contour, Detection, SegmentedObject, SegmentedObjectWithRects
 
 
 def add_rotated_rects(segmented_objects: list[SegmentedObject]) -> list[SegmentedObjectWithRects]:

model_api/python/model_api/models/yolo.py

@@ -11,7 +11,7 @@
 from model_api.adapters.utils import INTERPOLATION_TYPES, resize_image_ocv
 
 from .detection_model import DetectionModel
-from .result_types import Detection, DetectionResult
+from .result import Detection, DetectionResult
 from .types import BooleanValue, ListValue, NumericalValue
 from .utils import clip_detections, multiclass_nms, nms
 
@@ -744,8 +744,7 @@ def parameters(cls):
             {
                 "agnostic_nms": BooleanValue(
                     description=(
-                        "If True, the model is agnostic to the number of classes, "
-                        "and all classes are considered as one"
+                        "If True, the model is agnostic to the number of classes, and all classes are considered as one"
                     ),
                     default_value=False,
                 ),

model_api/python/model_api/visualizer/__init__.py

@@ -0,0 +1,10 @@
+"""Visualizer."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .media import Media
+from .primitives import Overlay
+from .visualizer import Visualizer
+
+__all__ = ["Media", "Visualizer"]

model_api/python/model_api/visualizer/layout.py

@@ -0,0 +1,40 @@
+"""Visualization Layout"""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from abc import ABC
+from typing import TYPE_CHECKING, Type
+
+if TYPE_CHECKING:
+    from PIL import Image
+
+    from model_api.visualizer.primitives import Primitive
+
+    from .media import Media
+
+
+class Layout(ABC):
+    """Base class for layouts."""
+
+    def _compute_on_primitive(self, primitive: Primitive, image: Image, media: Media) -> Image | None:
+        if media.has_primitive(primitive):
+            primitives = media.get_primitive(primitive)
+            for primitive in primitives:
+                image = primitive.compute(image)
+            return image
+        return None
+
+
+class Flatten(Layout):
+    """Put all primitives on top of each other"""
+
+    def __init__(self, *args: Type[Primitive]) -> None:
+        self.children = args
+
+    def __call__(self, image: Image, media: Media) -> Image:
+        _image: Image = image.copy()
+        for child in self.children:
+            _image = self._compute_on_primitive(child, _image, media)
+        return _image

model_api/python/model_api/visualizer/media.py

@@ -0,0 +1,58 @@
+"""Media object."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import abstractmethod
+from typing import Type
+
+from .layout import Layout
+from .primitives import Overlay, Primitive
+
+
+class Media:
+    """Media object.
+
+    Media object that is used by the visualizer to render the prediction.
+
+    Args:
+        *args: Primitives to be added to the prediction.
+
+    Example:
+        >>> media = Media(Label("Label"), BoundingBoxes(0, 0, 10, 10))
+    """
+
+    def __init__(self, *args: Primitive) -> None:
+        self._overlays: list[Overlay] = []
+        self._add_primitives(args)
+
+    def _add_primitives(self, primitives: list[Primitive]) -> None:
+        """Add primitives to the prediction."""
+        for primitive in primitives:
+            self._add_primitive(primitive)
+
+    def _add_primitive(self, primitive: Primitive) -> None:
+        """Add primitive."""
+        if isinstance(primitive, Overlay):
+            self._overlays.append(primitive)
+        else:
+            msg = f"Primitive {primitive} not supported"
+            raise ValueError(msg)
+
+    def has_primitive(self, primitive: Type[Primitive]) -> bool:
+        """Check if the primitive type is registered."""
+        if primitive == Overlay:
+            return bool(self._overlays)
+        return False
+
+    def get_primitive(self, primitive: Type[Primitive]) -> Primitive:
+        """Get primitive."""
+        if primitive == Overlay:
+            return self._overlays
+        msg = f"Primitive {primitive} not found"
+        raise ValueError(msg)
+
+    @property
+    @abstractmethod
+    def default_layout(self) -> Layout:
+        """Default layout for the media."""

model_api/python/model_api/visualizer/primitives.py

@@ -0,0 +1,44 @@
+"""Base class for primitives."""
+
+# Copyright (C) 2024 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from io import BytesIO
+from typing import TYPE_CHECKING
+
+import cv2
+from PIL import Image, ImageDraw, ImageFont
+
+if TYPE_CHECKING:
+    import numpy as np
+
+
+class Primitive(ABC):
+    """Primitive class."""
+
+    @abstractmethod
+    def compute(self, **kwargs) -> Image:
+        pass
+
+
+class Overlay(Primitive):
+    """Overlay an image.
+
+    Useful for XAI and Anomaly Maps.
+    """
+
+    def __init__(self, image: Image | np.ndarray, opacity: float = 0.4) -> None:
+        self.image = self._to_image(image)
+        self.opacity = opacity
+
+    def _to_image(self, image: Image | np.ndarray) -> Image:
+        if isinstance(image, Image.Image):
+            return image
+        return Image.fromarray(image)
+
+    def compute(self, image: Image) -> Image:
+        _image = self.image.resize(image.size)
+        return Image.blend(image, _image, self.opacity)