bug-fix + add classification

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

Author: ashwinvaidya17

model_api/python/model_api/models/result_types.py

@@ -30,6 +30,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
@@ -73,6 +74,7 @@ def __init__(
         feature_vector: np.ndarray | None = None,
         raw_scores: np.ndarray | None = None,
     ) -> None:
+        super().__init__()
         self.top_labels = top_labels
         self.saliency_map = saliency_map
         self.feature_vector = feature_vector
@@ -87,7 +89,9 @@ def __str__(self) -> str:
         )
 
     def _register_primitives(self) -> None:
-        pass
+        # TODO add saliency map
+        for idx, label, confidence in self.top_labels:
+            self._add_primitive(Label(f"Rank: {idx}, {label}: {confidence:.3f}"))
 
 
 class Detection:

model_api/python/model_api/visualizer/visualize_mixin.py

@@ -11,11 +11,12 @@
 class VisualizeMixin(ABC):
     """Mixin for visualization."""
 
-    _labels: list[Label] = []
-    _polygons: list[Polygon] = []
-    _overlays: list[Overlay] = []
-    _bounding_boxes: list[BoundingBoxes] = []
-    _registered_primitives: bool = False
+    def __init__(self) -> None:
+        self._labels = []
+        self._polygons = []
+        self._overlays = []
+        self._bounding_boxes = []
+        self._registered_primitives = False
 
     @abstractmethod
     def _register_primitives(self) -> None:

model_api/python/model_api/visualizer/visualizer.py

@@ -57,30 +57,31 @@ def _generate(self, image: Image, result: VisualizeMixin, visualization_type: Vi
     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)
+                image = overlay.compute(_image)
 
         elif result.has_polygons:  # 2. else use polygons
             polygons = result.get_polygons()
             for polygon in polygons:
-                image = polygon.compute(image)
+                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)
+                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)
+                label_images.append(label.compute(_image, overlay_on_image=False))
+            _image = Label.overlay_labels(_image, label_images)
 
-        return image
+        return _image
 
     def _generate_full(self, image: Image, result: VisualizeMixin) -> Image:
         """Return a single image with visualizations side by side."""