Add title to overlay

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

Author: ashwinvaidya17

src/python/model_api/visualizer/layout/hstack.py

@@ -9,6 +9,8 @@
 
 import PIL
 
+from model_api.visualizer.primitive import Overlay
+
 from .layout import Layout
 
 if TYPE_CHECKING:
@@ -31,6 +33,8 @@ def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, sc
             images = []
             for _primitive in scene.get_primitives(primitive):
                 image_ = _primitive.compute(image.copy())
+                if isinstance(_primitive, Overlay):
+                    image_ = Overlay.overlay_labels(image=image_, labels=_primitive.label)
                 images.append(image_)
             return self._stitch(*images)
         return None

src/python/model_api/visualizer/primitive/overlay.py

@@ -5,8 +5,11 @@
 
 from __future__ import annotations
 
+from typing import Union
+
 import numpy as np
 import PIL
+from PIL import ImageFont
 
 from .primitive import Primitive
 
@@ -18,11 +21,18 @@ class Overlay(Primitive):
 
     Args:
         image (PIL.Image | np.ndarray): Image to be overlaid.
+        label (str | None): Optional label name to overlay.
         opacity (float): Opacity of the overlay.
     """
 
-    def __init__(self, image: PIL.Image | np.ndarray, opacity: float = 0.4) -> None:
+    def __init__(
+        self,
+        image: PIL.Image | np.ndarray,
+        opacity: float = 0.4,
+        label: Union[str, None] = None,
+    ) -> None:
         self.image = self._to_pil(image)
+        self.label = label
         self.opacity = opacity
 
     def _to_pil(self, image: PIL.Image | np.ndarray) -> PIL.Image:
@@ -33,3 +43,22 @@ def _to_pil(self, image: PIL.Image | np.ndarray) -> PIL.Image:
     def compute(self, image: PIL.Image) -> PIL.Image:
         image_ = self.image.resize(image.size)
         return PIL.Image.blend(image, image_, self.opacity)
+
+    @classmethod
+    def overlay_labels(cls, image: PIL.Image, labels: Union[list[str], str, None] = None) -> PIL.Image:
+        """Draw labels at the bottom center of the image.
+
+        This is handy when you want to add a label to the image.
+        """
+        if labels is not None:
+            labels = [labels] if isinstance(labels, str) else labels
+            font = ImageFont.load_default(size=18)
+            buffer_y = 5
+            dummy_image = PIL.Image.new("RGB", (1, 1))
+            draw = PIL.ImageDraw.Draw(dummy_image)
+            textbox = draw.textbbox((0, 0), ", ".join(labels), font=font)
+            image_ = PIL.Image.new("RGB", (textbox[2] - textbox[0], textbox[3] + buffer_y - textbox[1]), "white")
+            draw = PIL.ImageDraw.Draw(image_)
+            draw.text((0, 0), ", ".join(labels), font=font, fill="black")
+            image.paste(image_, (image.width // 2 - image_.width // 2, image.height - image_.height - buffer_y))
+        return image

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

@@ -3,7 +3,6 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from itertools import starmap
 from typing import Union
 
 import cv2
@@ -22,27 +21,28 @@ class DetectionScene(Scene):
     def __init__(self, image: Image, result: DetectionResult, layout: Union[Layout, None] = None) -> None:
         super().__init__(
             base=image,
-            label=self._get_labels(result),
             bounding_box=self._get_bounding_boxes(result),
             overlay=self._get_overlays(result),
             layout=layout,
         )
 
-    def _get_labels(self, result: DetectionResult) -> list[Label]:
-        labels = []
-        for label, score, label_name in zip(result.labels, result.scores, result.label_names):
-            labels.append(Label(label=f"{label} {label_name}", score=score))
-        return labels
-
     def _get_overlays(self, result: DetectionResult) -> list[Overlay]:
         overlays = []
-        for saliency_map in result.saliency_map[0][1:]:  # Assumes only one batch. Skip background class.
-            saliency_map = cv2.applyColorMap(saliency_map, cv2.COLORMAP_JET)
-            overlays.append(Overlay(saliency_map))
+        # Add only the overlays that are predicted
+        label_index_mapping = dict(zip(result.labels, result.label_names))
+        for label_index, label_name in label_index_mapping.items():
+            # Index 0 as it assumes only one batch
+            saliency_map = cv2.applyColorMap(result.saliency_map[0][label_index], cv2.COLORMAP_JET)
+            overlays.append(Overlay(saliency_map, label=label_name.title()))
         return overlays
 
     def _get_bounding_boxes(self, result: DetectionResult) -> list[BoundingBox]:
-        return list(starmap(BoundingBox, result.bboxes))
+        bounding_boxes = []
+        for score, label_name, bbox in zip(result.scores, result.label_names, result.bboxes):
+            x1, y1, x2, y2 = bbox
+            label = f"{label_name} ({score:.2f})"
+            bounding_boxes.append(BoundingBox(x1=x1, y1=y1, x2=x2, y2=y2, label=label))
+        return bounding_boxes
 
     @property
     def default_layout(self) -> Layout: