Add classification scene

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

Author: ashwinvaidya17

src/python/model_api/visualizer/scene/classification.py

@@ -5,11 +5,12 @@
 
 from typing import Union
 
+import cv2
 from PIL import Image
 
 from model_api.models.result import ClassificationResult
 from model_api.visualizer.layout import Flatten, Layout
-from model_api.visualizer.primitive import Overlay
+from model_api.visualizer.primitive import Label, Overlay
 
 from .scene import Scene
 
@@ -18,9 +19,26 @@ class ClassificationScene(Scene):
     """Classification Scene."""
 
     def __init__(self, image: Image, result: ClassificationResult, layout: Union[Layout, None] = None) -> None:
-        self.image = image
-        self.result = result
+        super().__init__(
+            base=image,
+            label=self._get_labels(result),
+            overlay=self._get_overlays(result),
+            layout=layout,
+        )
+
+    def _get_labels(self, result: ClassificationResult) -> list[Label]:
+        labels = []
+        if result.top_labels is not None and len(result.top_labels) > 0:
+            labels.extend([Label(label=str(label)) for label in result.top_labels])
+        return labels
+
+    def _get_overlays(self, result: ClassificationResult) -> list[Overlay]:
+        overlays = []
+        if result.saliency_map is not None and result.saliency_map.size > 0:
+            saliency_map = cv2.cvtColor(result.saliency_map, cv2.COLOR_BGR2RGB)
+            overlays.append(Overlay(saliency_map))
+        return overlays
 
     @property
     def default_layout(self) -> Layout:
-        return Flatten(Overlay)
+        return Flatten(Overlay, Label)

tests/python/unit/visualizer/test_scene.py

@@ -8,7 +8,8 @@
 import numpy as np
 from PIL import Image
 
-from model_api.models.result import AnomalyResult
+from model_api.models.result import AnomalyResult, ClassificationResult
+from model_api.models.result.classification import Label
 from model_api.visualizer import Visualizer
 
 
@@ -32,3 +33,19 @@ def test_anomaly_scene(mock_image: Image, tmpdir: Path):
     visualizer = Visualizer()
     visualizer.save(mock_image, anomaly_result, tmpdir / "anomaly_scene.jpg")
     assert Path(tmpdir / "anomaly_scene.jpg").exists()
+
+
+def test_classification_scene(mock_image: Image, tmpdir: Path):
+    """Test if the classification scene is created."""
+    classification_result = ClassificationResult(
+        top_labels=[
+            Label(name="cat", confidence=0.95),
+            Label(name="dog", confidence=0.90),
+        ],
+        saliency_map=np.ones(mock_image.size, dtype=np.uint8),
+    )
+    visualizer = Visualizer()
+    visualizer.save(
+        mock_image, classification_result, tmpdir / "classification_scene.jpg"
+    )
+    assert Path(tmpdir / "classification_scene.jpg").exists()