From d520767057a67556db74d43369694bae2bcb17dc Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Wed, 22 Jan 2025 13:42:10 +0100 Subject: [PATCH 1/6] Add classification scene Signed-off-by: Ashwin Vaidya --- .../visualizer/scene/classification.py | 26 ++++++++++++++++--- tests/python/unit/visualizer/test_scene.py | 19 +++++++++++++- 2 files changed, 40 insertions(+), 5 deletions(-) diff --git a/src/python/model_api/visualizer/scene/classification.py b/src/python/model_api/visualizer/scene/classification.py index ec2f3fe6..da37d7ca 100644 --- a/src/python/model_api/visualizer/scene/classification.py +++ b/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) diff --git a/tests/python/unit/visualizer/test_scene.py b/tests/python/unit/visualizer/test_scene.py index b25538e5..a2a9ba7f 100644 --- a/tests/python/unit/visualizer/test_scene.py +++ b/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() From 8bf9708ca35fec3fc212bb1d24ce7d3c8cdd0aa7 Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Wed, 22 Jan 2025 14:12:09 +0100 Subject: [PATCH 2/6] Add detection scene Signed-off-by: Ashwin Vaidya --- .../model_api/models/result/detection.py | 5 +++ .../model_api/visualizer/scene/detection.py | 34 +++++++++++++++++-- 2 files changed, 36 insertions(+), 3 deletions(-) diff --git a/src/python/model_api/models/result/detection.py b/src/python/model_api/models/result/detection.py index 40df2b0e..f461f610 100644 --- a/src/python/model_api/models/result/detection.py +++ b/src/python/model_api/models/result/detection.py @@ -111,6 +111,11 @@ def label_names(self, value): @property def saliency_map(self): + """Saliency map for XAI. + + Returns: + np.ndarray: Saliency map in dim of (B, N_CLASSES, H, W). + """ return self._saliency_map @saliency_map.setter diff --git a/src/python/model_api/visualizer/scene/detection.py b/src/python/model_api/visualizer/scene/detection.py index 2cc7691a..60c823ca 100644 --- a/src/python/model_api/visualizer/scene/detection.py +++ b/src/python/model_api/visualizer/scene/detection.py @@ -3,12 +3,15 @@ # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 +from itertools import starmap from typing import Union +import cv2 from PIL import Image from model_api.models.result import DetectionResult -from model_api.visualizer.layout import Layout +from model_api.visualizer.layout import Flatten, HStack, Layout +from model_api.visualizer.primitive import BoundingBox, Label, Overlay from .scene import Scene @@ -17,5 +20,30 @@ class DetectionScene(Scene): """Detection Scene.""" def __init__(self, image: Image, result: DetectionResult, layout: Union[Layout, None] = None) -> None: - self.image = image - self.result = result + 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)) + return overlays + + def _get_bounding_boxes(self, result: DetectionResult) -> list[BoundingBox]: + return list(starmap(BoundingBox, result.bboxes)) + + @property + def default_layout(self) -> Layout: + return HStack(Flatten(BoundingBox, Label), Overlay) From e59de98693c247c75c08cabd26e790ab124b7aca Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Wed, 22 Jan 2025 14:16:41 +0100 Subject: [PATCH 3/6] Add tests Signed-off-by: Ashwin Vaidya --- tests/python/unit/visualizer/test_scene.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/python/unit/visualizer/test_scene.py b/tests/python/unit/visualizer/test_scene.py index a2a9ba7f..cee09399 100644 --- a/tests/python/unit/visualizer/test_scene.py +++ b/tests/python/unit/visualizer/test_scene.py @@ -8,7 +8,7 @@ import numpy as np from PIL import Image -from model_api.models.result import AnomalyResult, ClassificationResult +from model_api.models.result import AnomalyResult, ClassificationResult, DetectionResult from model_api.models.result.classification import Label from model_api.visualizer import Visualizer @@ -49,3 +49,17 @@ def test_classification_scene(mock_image: Image, tmpdir: Path): mock_image, classification_result, tmpdir / "classification_scene.jpg" ) assert Path(tmpdir / "classification_scene.jpg").exists() + + +def test_detection_scene(mock_image: Image, tmpdir: Path): + """Test if the detection scene is created.""" + detection_result = DetectionResult( + bboxes=np.array([[0, 0, 128, 128], [32, 32, 96, 96]]), + labels=np.array([1, 2]), + label_names=["person", "car"], + scores=np.array([0.85, 0.75]), + saliency_map=(np.ones((1, 2, 6, 8)) * 255).astype(np.uint8), + ) + visualizer = Visualizer() + visualizer.save(mock_image, detection_result, tmpdir / "detection_scene.jpg") + assert Path(tmpdir / "detection_scene.jpg").exists() From ddcc269c833295b3da756e94f4493e085b57fc5e Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Wed, 22 Jan 2025 15:24:19 +0100 Subject: [PATCH 4/6] Add title to overlay Signed-off-by: Ashwin Vaidya --- .../model_api/visualizer/layout/hstack.py | 4 +++ .../model_api/visualizer/primitive/overlay.py | 31 ++++++++++++++++++- .../model_api/visualizer/scene/detection.py | 24 +++++++------- 3 files changed, 46 insertions(+), 13 deletions(-) diff --git a/src/python/model_api/visualizer/layout/hstack.py b/src/python/model_api/visualizer/layout/hstack.py index 6eb9e87a..2cc4079a 100644 --- a/src/python/model_api/visualizer/layout/hstack.py +++ b/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 diff --git a/src/python/model_api/visualizer/primitive/overlay.py b/src/python/model_api/visualizer/primitive/overlay.py index f69bc958..870e14de 100644 --- a/src/python/model_api/visualizer/primitive/overlay.py +++ b/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 diff --git a/src/python/model_api/visualizer/scene/detection.py b/src/python/model_api/visualizer/scene/detection.py index 60c823ca..aee838c3 100644 --- a/src/python/model_api/visualizer/scene/detection.py +++ b/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: From 36a6c561e0e578ab5531dccd75624873af56d2c3 Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Thu, 23 Jan 2025 09:20:23 +0100 Subject: [PATCH 5/6] Pass name and confidence separately Signed-off-by: Ashwin Vaidya --- src/python/model_api/visualizer/scene/classification.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/python/model_api/visualizer/scene/classification.py b/src/python/model_api/visualizer/scene/classification.py index da37d7ca..04d08774 100644 --- a/src/python/model_api/visualizer/scene/classification.py +++ b/src/python/model_api/visualizer/scene/classification.py @@ -29,7 +29,9 @@ def __init__(self, image: Image, result: ClassificationResult, layout: Union[Lay 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]) + for label in result.top_labels: + if label.name is not None: + labels.append(Label(label=label.name, score=label.confidence)) return labels def _get_overlays(self, result: ClassificationResult) -> list[Overlay]: From 13b763701bf278665724d3ce3b305e97d3c14d60 Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Thu, 23 Jan 2025 09:50:28 +0100 Subject: [PATCH 6/6] Fix tests Signed-off-by: Ashwin Vaidya --- tests/python/unit/visualizer/test_scene.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/unit/visualizer/test_scene.py b/tests/python/unit/visualizer/test_scene.py index cee09399..7b560443 100644 --- a/tests/python/unit/visualizer/test_scene.py +++ b/tests/python/unit/visualizer/test_scene.py @@ -55,7 +55,7 @@ def test_detection_scene(mock_image: Image, tmpdir: Path): """Test if the detection scene is created.""" detection_result = DetectionResult( bboxes=np.array([[0, 0, 128, 128], [32, 32, 96, 96]]), - labels=np.array([1, 2]), + labels=np.array([0, 1]), label_names=["person", "car"], scores=np.array([0.85, 0.75]), saliency_map=(np.ones((1, 2, 6, 8)) * 255).astype(np.uint8),