Add PoC 2

Signed-off-by: Ashwin Vaidya <[email protected]>

Author: ashwinvaidya17

model_api/python/model_api/models/result_types/anomaly.py

@@ -8,6 +8,7 @@
 import cv2
 import numpy as np
 
+from model_api.visualizer.layout import Flatten, Layout
 from model_api.visualizer.primitives import BoundingBoxes, Label, Overlay, Polygon
 
 from .base import Result
@@ -57,3 +58,11 @@ def _register_primitives(self) -> None:
             self._add_primitive(Label(self.pred_label, bg_color="red" if self.pred_label == "Anomaly" else "green"))
         self._add_primitive(Label(f"Score: {self.pred_score}"))
         self._add_primitive(Polygon(mask=self.pred_mask))
+
+    @property
+    def default_layout(self) -> Layout:
+        return Flatten(
+            Overlay,
+            Polygon,
+            Label,
+        )

model_api/python/model_api/visualizer/layout.py

@@ -0,0 +1,85 @@
+"""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
+
+from PIL import Image
+
+if TYPE_CHECKING:
+    from model_api.visualizer.primitives import Primitive
+
+    from .visualize_mixin import VisualizeMixin
+
+
+class Layout(ABC):
+    """Base class for layouts."""
+
+    def _compute_on_primitive(self, primitive: Primitive, image: Image, result: VisualizeMixin) -> Image | None:
+        if result.has_primitive(primitive):
+            primitives = result.get_primitive(primitive)
+            for primitive in primitives:
+                image = primitive.compute(image)
+            return image
+        return None
+
+
+class HStack(Layout):
+    """Horizontal stack layout."""
+
+    def __init__(self, *args: Layout | Type[Primitive]) -> None:
+        self.children = args
+
+    def __call__(self, image: Image, result: VisualizeMixin) -> Image:
+        images: list[Image] = []
+        for child in self.children:
+            if isinstance(child, Layout):
+                images.append(child(image, result))
+            else:
+                _image = image.copy()
+                _image = self._compute_on_primitive(child, _image, result)
+                if _image is not None:
+                    images.append(_image)
+        return self._stitch(*images)
+
+    def _stitch(self, *images: Image) -> Image:
+        """Stitch images together.
+
+        Args:
+            images (Image): Images to stitch.
+
+        Returns:
+            Image: Stitched image.
+        """
+        new_image = Image.new(
+            "RGB",
+            (
+                sum(image.width for image in images),
+                max(image.height for image in images),
+            ),
+        )
+        x_offset = 0
+        for image in images:
+            new_image.paste(image, (x_offset, 0))
+            x_offset += image.width
+        return new_image
+
+
+class VStack(Layout):
+    """Vertical stack layout."""
+
+
+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, result: VisualizeMixin) -> Image:
+        _image: Image = image.copy()
+        for child in self.children:
+            _image = self._compute_on_primitive(child, _image, result)
+        return _image

model_api/python/model_api/visualizer/visualize_mixin.py

@@ -4,7 +4,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
+from typing import Type
 
+from .layout import Layout
 from .primitives import BoundingBoxes, Label, Overlay, Polygon, Primitive
 
 
@@ -22,6 +24,11 @@ def __init__(self) -> None:
     def _register_primitives(self) -> None:
         """Convert result entities to primitives."""
 
+    @property
+    @abstractmethod
+    def default_layout(self) -> Layout:
+        """Default layout."""
+
     def _add_primitive(self, primitive: Primitive) -> None:
         """Add primitive."""
         if isinstance(primitive, Label):
@@ -33,49 +40,32 @@ def _add_primitive(self, primitive: Primitive) -> None:
         elif isinstance(primitive, BoundingBoxes):
             self._bounding_boxes.append(primitive)
 
-    @property
-    def has_labels(self) -> bool:
-        """Check if there are labels."""
-        self._register_primitives_if_needed()
-        return bool(self._labels)
-
-    @property
-    def has_bounding_boxes(self) -> bool:
-        """Check if there are bounding boxes."""
-        self._register_primitives_if_needed()
-        return bool(self._bounding_boxes)
-
-    @property
-    def has_polygons(self) -> bool:
-        """Check if there are polygons."""
-        self._register_primitives_if_needed()
-        return bool(self._polygons)
-
-    @property
-    def has_overlays(self) -> bool:
-        """Check if there are overlays."""
-        self._register_primitives_if_needed()
-        return bool(self._overlays)
-
-    def get_labels(self) -> list[Label]:
-        """Get labels."""
+    def has_primitive(self, primitive: Type[Primitive]) -> bool:
+        """Check if the primitive type is registered."""
         self._register_primitives_if_needed()
-        return self._labels
-
-    def get_polygons(self) -> list[Polygon]:
-        """Get polygons."""
-        self._register_primitives_if_needed()
-        return self._polygons
-
-    def get_overlays(self) -> list[Overlay]:
-        """Get overlays."""
-        self._register_primitives_if_needed()
-        return self._overlays
-
-    def get_bounding_boxes(self) -> list[BoundingBoxes]:
-        """Get bounding boxes."""
+        if primitive == Label:
+            return bool(self._labels)
+        if primitive == Polygon:
+            return bool(self._polygons)
+        if primitive == Overlay:
+            return bool(self._overlays)
+        if primitive == BoundingBoxes:
+            return bool(self._bounding_boxes)
+        return False
+
+    def get_primitive(self, primitive: Type[Primitive]) -> Primitive:
+        """Get primitive."""
         self._register_primitives_if_needed()
-        return self._bounding_boxes
+        if primitive == Label:
+            return self._labels
+        if primitive == Polygon:
+            return self._polygons
+        if primitive == Overlay:
+            return self._overlays
+        if primitive == BoundingBoxes:
+            return self._bounding_boxes
+        msg = f"Primitive {primitive} not found"
+        raise ValueError(msg)
 
     def _register_primitives_if_needed(self):
         if not self._registered_primitives:

model_api/python/model_api/visualizer/visualizer.py

@@ -8,130 +8,38 @@
 from enum import Enum
 from typing import TYPE_CHECKING
 
-from PIL import Image
-
 from model_api.visualizer.primitives import Label
 
 if TYPE_CHECKING:
-    from model_api.visualizer.visualize_mixin import VisualizeMixin
+    from PIL import Image
 
+    from model_api.visualizer.visualize_mixin import VisualizeMixin
 
-class VisualizationType(Enum):
-    """Visualization type."""
-
-    FULL = "full"
-    SIMPLE = "simple"
+    from .layout import Layout
 
 
 class Visualizer:
-    def __init__(self) -> None:
-        # TODO: add transforms for the source image so that it has the same crop, and size as the model.
-        pass
+    def __init__(self, layout: Layout | None = None) -> None:
+        self.layout = layout
 
     def show(
         self,
         image: Image,
         result: VisualizeMixin,
-        visualization_type: VisualizationType | str = VisualizationType.SIMPLE,
     ) -> None:
-        visualization_type = VisualizationType(visualization_type)
-        result: Image = self._generate(image, result, visualization_type)
+        result: Image = self._generate(image, result)
         result.show()
 
     def save(
         self,
         image: Image,
         result: VisualizeMixin,
         path: str,
-        visualization_type: VisualizationType | str = VisualizationType.SIMPLE,
     ) -> None:
-        visualization_type = VisualizationType(visualization_type)
-        result: Image = self._generate(image, result, visualization_type)
+        result: Image = self._generate(image, result)
         result.save(path)
 
-    def _generate(self, image: Image, result: VisualizeMixin, visualization_type: VisualizationType) -> Image:
-        _result: Image
-        if visualization_type == VisualizationType.SIMPLE:
-            _result = self._generate_simple(image, result)
-        else:
-            _result = self._generate_full(image, result)
-        return _result
-
-    def _generate_simple(self, image: Image, result: VisualizeMixin) -> Image:
-        """Return a single image with stacked visualizations."""
-        # 1. Use Overlay
-        _image = image.copy()
-        if result.has_overlays:
-            overlays = result.get_overlays()
-            for overlay in overlays:
-                image = overlay.compute(_image)
-
-        elif result.has_polygons:  # 2. else use polygons
-            polygons = result.get_polygons()
-            for polygon in polygons:
-                image = polygon.compute(_image)
-
-        elif result.has_bounding_boxes:  # 3. else use bounding boxes
-            bounding_boxes = result.get_bounding_boxes()
-            for bounding_box in bounding_boxes:
-                image = bounding_box.compute(_image)
-
-        # Finally add labels
-        if result.has_labels:
-            labels = result.get_labels()
-            label_images = []
-            for label in labels:
-                label_images.append(label.compute(_image, overlay_on_image=False))
-            _image = Label.overlay_labels(_image, label_images)
-
-        return _image
-
-    def _generate_full(self, image: Image, result: VisualizeMixin) -> Image:
-        """Return a single image with visualizations side by side."""
-        images: list[Image] = [image]
-
-        if result.has_overlays:
-            overlays = result.get_overlays()
-            _image = image.copy()
-            for overlay in overlays:
-                _image = overlay.compute(_image)
-            images.append(_image)
-        if result.has_polygons:
-            polygons = result.get_polygons()
-            _image = image.copy()
-            for polygon in polygons:
-                _image = polygon.compute(_image)
-            images.append(_image)
-        if result.has_bounding_boxes:
-            bounding_boxes = result.get_bounding_boxes()
-            _image = image.copy()
-            for bounding_box in bounding_boxes:
-                _image = bounding_box.compute(_image)
-            images.append(_image)
-        if result.has_labels:
-            labels = result.get_labels()
-            for label in labels:
-                images.append(label.compute(image.copy(), overlay_on_image=True))
-        return self._stitch(*images)
-
-    def _stitch(self, *images: Image) -> Image:
-        """Stitch images together.
-
-        Args:
-            images (Image): Images to stitch.
-
-        Returns:
-            Image: Stitched image.
-        """
-        new_image = Image.new(
-            "RGB",
-            (
-                sum(image.width for image in images),
-                max(image.height for image in images),
-            ),
-        )
-        x_offset = 0
-        for image in images:
-            new_image.paste(image, (x_offset, 0))
-            x_offset += image.width
-        return new_image
+    def _generate(self, image: Image, result: VisualizeMixin) -> Image:
+        if self.layout is not None:
+            return self.layout(image, result)
+        return result.default_layout(image, result)