restore visualization changes

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

Author: ashwinvaidya17

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.primitives import BoundingBoxes, Label, Overlay, Polygon
 
-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
@@ -40,3 +46,14 @@ def __str__(self) -> str:
             f"pred_label:{self.pred_label};"
             f"pred_mask min:{pred_mask_min} max:{pred_mask_max};"
         )
+
+    def _register_primitives(self) -> None:
+        """Converts the result to primitives."""
+        anomaly_map = cv2.applyColorMap(self.anomaly_map, cv2.COLORMAP_JET)
+        self._add_primitive(Overlay(anomaly_map))
+        for box in self.pred_boxes:
+            self._add_primitive(BoundingBoxes(*box))
+        if self.pred_label is not 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))

model_api/python/model_api/models/result_types/base.py

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

model_api/python/model_api/models/result_types/classification.py

@@ -7,13 +7,16 @@
 
 from typing import TYPE_CHECKING
 
+from model_api.visualizer.primitives import Label
+
+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__(
@@ -35,3 +38,8 @@ def __str__(self) -> str:
             f"{labels}, {array_shape_to_str(self.saliency_map)}, {array_shape_to_str(self.feature_vector)}, "
             f"{array_shape_to_str(self.raw_scores)}"
         )
+
+    def _register_primitives(self) -> None:
+        # TODO add saliency map
+        for idx, label, confidence in self.top_labels:
+            self._add_primitive(Label(f"Rank: {idx}, {label}: {confidence:.3f}"))

model_api/python/model_api/visualizer/visualizer.py

@@ -6,11 +6,14 @@
 from __future__ import annotations
 
 from enum import Enum
+from typing import TYPE_CHECKING
 
 from PIL import Image
 
 from model_api.visualizer.primitives import Label
-from model_api.visualizer.visualize_mixin import VisualizeMixin
+
+if TYPE_CHECKING:
+    from model_api.visualizer.visualize_mixin import VisualizeMixin
 
 
 class VisualizationType(Enum):
@@ -47,12 +50,12 @@ def save(
         result.save(path)
 
     def _generate(self, image: Image, result: VisualizeMixin, visualization_type: VisualizationType) -> Image:
-        result: Image
+        _result: Image
         if visualization_type == VisualizationType.SIMPLE:
-            result = self._generate_simple(image, result)
+            _result = self._generate_simple(image, result)
         else:
-            result = self._generate_full(image, result)
-        return result
+            _result = self._generate_full(image, result)
+        return _result
 
     def _generate_simple(self, image: Image, result: VisualizeMixin) -> Image:
         """Return a single image with stacked visualizations."""