Add Keypoint and Instance Segmentation support in visualizer (#261)

* Add Keypoint and Instance Segmentation support in visualizer

- Introduced Keypoint primitive for visualizing keypoints on images.
- Updated visualizer to handle Instance Segmentation results and added corresponding scene.
- Refactored existing scenes to integrate new functionalities, including KeypointScene and InstanceSegmentationScene.
- Enhanced Polygon primitive with opacity and outline width options.
- Updated tests to cover new features and ensure functionality.

This commit enhances the visualizer's capabilities for handling keypoints and instance segmentation, improving the overall visualization experience.

* fix saliency map colour conversion

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

* Add visualization example for VisionAPI

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

---------

Signed-off-by: Ashwin Vaidya <ashwinnitinvaidya@gmail.com>

Author: ashwinvaidya17

examples/python/visualization/README.md

@@ -0,0 +1,15 @@
+# Visualization Example
+
+This example demonstrates how to use the Visualizer in VisionAPI.
+
+## Prerequisites
+
+Install Model API from source. Please refer to the main [README](../../../README.md) for details.
+
+## Run example
+
+To run the example, please execute the following command:
+
+```bash
+python run.py --image <path_to_image> --model <path_to_model>.xml --output <path_to_output_image>
+```

examples/python/visualization/run.py

@@ -0,0 +1,38 @@
+"""Visualization Example."""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import argparse
+from argparse import Namespace
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from model_api.models import Model
+from model_api.visualizer import Visualizer
+
+
+def main(args: Namespace):
+    image = Image.open(args.image)
+
+    model = Model.create_model(args.model)
+
+    image_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    predictions = model(image_array)
+    visualizer = Visualizer()
+
+    if args.output:
+        visualizer.save(image=image, result=predictions, path=args.output)
+    else:
+        visualizer.show(image=image, result=predictions)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image", type=str, required=True)
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument("--output", type=str, required=False)
+    args = parser.parse_args()
+    main(args)

src/python/model_api/visualizer/__init__.py

@@ -4,8 +4,19 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .layout import Flatten, HStack, Layout
-from .primitive import BoundingBox, Label, Overlay, Polygon
+from .primitive import BoundingBox, Keypoint, Label, Overlay, Polygon
 from .scene import Scene
 from .visualizer import Visualizer
 
-__all__ = ["BoundingBox", "Label", "Overlay", "Polygon", "Scene", "Visualizer", "Layout", "Flatten", "HStack"]
+__all__ = [
+    "BoundingBox",
+    "Keypoint",
+    "Label",
+    "Overlay",
+    "Polygon",
+    "Scene",
+    "Visualizer",
+    "Layout",
+    "Flatten",
+    "HStack",
+]

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

@@ -4,9 +4,10 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .bounding_box import BoundingBox
+from .keypoints import Keypoint
 from .label import Label
 from .overlay import Overlay
 from .polygon import Polygon
 from .primitive import Primitive
 
-__all__ = ["Primitive", "BoundingBox", "Label", "Overlay", "Polygon"]
+__all__ = ["Primitive", "BoundingBox", "Label", "Overlay", "Polygon", "Keypoint"]

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

@@ -0,0 +1,65 @@
+"""Keypoints primitive."""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Union
+
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+
+from .primitive import Primitive
+
+
+class Keypoint(Primitive):
+    """Keypoint primitive.
+
+    Args:
+        keypoints (np.ndarray): Keypoints. Shape: (N, 2)
+        scores (np.ndarray | None): Scores. Shape: (N,). Defaults to None.
+        color (str | tuple[int, int, int]): Color of the keypoints. Defaults to "purple".
+    """
+
+    def __init__(
+        self,
+        keypoints: np.ndarray,
+        scores: Union[np.ndarray, None] = None,
+        color: Union[str, tuple[int, int, int]] = "purple",
+        keypoint_size: int = 3,
+    ) -> None:
+        self.keypoints = self._validate_keypoints(keypoints)
+        self.scores = scores
+        self.color = color
+        self.keypoint_size = keypoint_size
+
+    def compute(self, image: Image) -> Image:
+        """Draw keypoints on the image."""
+        draw = ImageDraw.Draw(image)
+        for keypoint in self.keypoints:
+            draw.ellipse(
+                (
+                    keypoint[0] - self.keypoint_size,
+                    keypoint[1] - self.keypoint_size,
+                    keypoint[0] + self.keypoint_size,
+                    keypoint[1] + self.keypoint_size,
+                ),
+                fill=self.color,
+            )
+
+        if self.scores is not None:
+            font = ImageFont.load_default(size=18)
+            for score, keypoint in zip(self.scores, self.keypoints):
+                textbox = draw.textbbox((0, 0), f"{score:.2f}", font=font)
+                draw.text(
+                    (keypoint[0] - textbox[2] // 2, keypoint[1] + self.keypoint_size),
+                    f"{score:.2f}",
+                    font=font,
+                    fill=self.color,
+                )
+        return image
+
+    def _validate_keypoints(self, keypoints: np.ndarray) -> np.ndarray:
+        if keypoints.shape[1] != 2:
+            msg = "Keypoints must have shape (N, 2)"
+            raise ValueError(msg)
+        return keypoints

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

@@ -5,16 +5,19 @@
 
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING
 
 import cv2
-from PIL import Image, ImageDraw
+from PIL import Image, ImageColor, ImageDraw
 
 from .primitive import Primitive
 
 if TYPE_CHECKING:
     import numpy as np
 
+logger = logging.getLogger(__name__)
+
 
 class Polygon(Primitive):
     """Polygon primitive.
@@ -38,9 +41,13 @@ def __init__(
         points: list[tuple[int, int]] | None = None,
         mask: np.ndarray | None = None,
         color: str | tuple[int, int, int] = "blue",
+        opacity: float = 0.4,
+        outline_width: int = 2,
     ) -> None:
         self.points = self._get_points(points, mask)
         self.color = color
+        self.opacity = opacity
+        self.outline_width = outline_width
 
     def _get_points(self, points: list[tuple[int, int]] | None, mask: np.ndarray | None) -> list[tuple[int, int]]:
         """Get points from either points or mask.
@@ -76,6 +83,13 @@ def _get_points_from_mask(self, mask: np.ndarray) -> list[tuple[int, int]]:
             List of points.
         """
         contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        # incase of multiple contours, use the one with the largest area
+        if len(contours) > 1:
+            logger.warning("Multiple contours found in the mask. Using the largest one.")
+            contours = sorted(contours, key=cv2.contourArea, reverse=True)
+        if len(contours) == 0:
+            msg = "No contours found in the mask."
+            raise ValueError(msg)
         points_ = contours[0].squeeze().tolist()
         return [tuple(point) for point in points_]
 
@@ -88,6 +102,8 @@ def compute(self, image: Image) -> Image:
         Returns:
             Image with the polygon drawn on it.
         """
-        draw = ImageDraw.Draw(image)
-        draw.polygon(self.points, fill=self.color)
+        draw = ImageDraw.Draw(image, "RGBA")
+        # Draw polygon with darker edge and a semi-transparent fill.
+        ink = ImageColor.getrgb(self.color)
+        draw.polygon(self.points, fill=(*ink, int(255 * self.opacity)), outline=self.color, width=self.outline_width)
         return image

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

@@ -8,13 +8,14 @@
 from .detection import DetectionScene
 from .keypoint import KeypointScene
 from .scene import Scene
-from .segmentation import SegmentationScene
+from .segmentation import InstanceSegmentationScene, SegmentationScene
 from .visual_prompting import VisualPromptingScene
 
 __all__ = [
     "AnomalyScene",
     "ClassificationScene",
     "DetectionScene",
+    "InstanceSegmentationScene",
     "KeypointScene",
     "Scene",
     "SegmentationScene",

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

@@ -32,8 +32,10 @@ def _get_overlays(self, result: DetectionResult) -> list[Overlay]:
         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()))
+            if result.saliency_map is not None and result.saliency_map.size > 0:
+                saliency_map = cv2.applyColorMap(result.saliency_map[0][label_index], cv2.COLORMAP_JET)
+                saliency_map = cv2.cvtColor(saliency_map, cv2.COLOR_BGR2RGB)
+                overlays.append(Overlay(saliency_map, label=label_name.title()))
         return overlays
 
     def _get_bounding_boxes(self, result: DetectionResult) -> list[BoundingBox]:

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

@@ -3,19 +3,30 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Union
+
+from PIL import Image
+
 from model_api.models.result import DetectedKeypoints
 from model_api.visualizer.layout import Flatten, Layout
-from model_api.visualizer.primitive import Overlay
+from model_api.visualizer.primitive import Keypoint
 
 from .scene import Scene
 
 
 class KeypointScene(Scene):
     """Keypoint Scene."""
 
-    def __init__(self, result: DetectedKeypoints) -> None:
-        self.result = result
+    def __init__(self, image: Image, result: DetectedKeypoints, layout: Union[Layout, None] = None) -> None:
+        super().__init__(
+            base=image,
+            keypoints=self._get_keypoints(result),
+            layout=layout,
+        )
+
+    def _get_keypoints(self, result: DetectedKeypoints) -> list[Keypoint]:
+        return [Keypoint(result.keypoints, result.scores)]
 
     @property
     def default_layout(self) -> Layout:
-        return Flatten(Overlay)
+        return Flatten(Keypoint)

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

@@ -10,7 +10,7 @@
 import numpy as np
 from PIL import Image
 
-from model_api.visualizer.primitive import BoundingBox, Label, Overlay, Polygon, Primitive
+from model_api.visualizer.primitive import BoundingBox, Keypoint, Label, Overlay, Polygon, Primitive
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -31,13 +31,15 @@ def __init__(
         label: Label | list[Label] | None = None,
         overlay: Overlay | list[Overlay] | np.ndarray | None = None,
         polygon: Polygon | list[Polygon] | None = None,
+        keypoints: Keypoint | list[Keypoint] | np.ndarray | None = None,
         layout: Layout | None = None,
     ) -> None:
         self.base = base
         self.overlay = self._to_overlay(overlay)
         self.bounding_box = self._to_bounding_box(bounding_box)
         self.label = self._to_label(label)
         self.polygon = self._to_polygon(polygon)
+        self.keypoints = self._to_keypoints(keypoints)
         self.layout = layout
 
     def show(self) -> None:
@@ -60,6 +62,8 @@ def has_primitives(self, primitive: type[Primitive]) -> bool:
             return bool(self.label)
         if primitive == Polygon:
             return bool(self.polygon)
+        if primitive == Keypoint:
+            return bool(self.keypoints)
         return False
 
     def get_primitives(self, primitive: type[Primitive]) -> list[Primitive]:
@@ -86,6 +90,8 @@ def get_primitives(self, primitive: type[Primitive]) -> list[Primitive]:
             primitives = cast("list[Primitive]", self.label)
         elif primitive == Polygon:
             primitives = cast("list[Primitive]", self.polygon)
+        elif primitive == Keypoint:
+            primitives = cast("list[Primitive]", self.keypoints)
         else:
             msg = f"Primitive {primitive} not found"
             raise ValueError(msg)
@@ -119,3 +125,10 @@ def _to_polygon(self, polygon: Polygon | list[Polygon] | None) -> list[Polygon]
         if isinstance(polygon, Polygon):
             return [polygon]
         return polygon
+
+    def _to_keypoints(self, keypoints: Keypoint | list[Keypoint] | np.ndarray | None) -> list[Keypoint] | None:
+        if isinstance(keypoints, Keypoint):
+            return [keypoints]
+        if isinstance(keypoints, np.ndarray):
+            return [Keypoint(keypoints)]
+        return keypoints

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

@@ -0,0 +1,12 @@
+"""Segmentation Scene."""
+
+# Copyright (C) 2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .instance_segmentation import InstanceSegmentationScene
+from .segmentation import SegmentationScene
+
+__all__ = [
+    "InstanceSegmentationScene",
+    "SegmentationScene",
+]