diff --git a/model_api/python/model_api/models/__init__.py b/model_api/python/model_api/models/__init__.py index eba5deeb..2b329c45 100644 --- a/model_api/python/model_api/models/__init__.py +++ b/model_api/python/model_api/models/__init__.py @@ -56,44 +56,47 @@ __all__ = [ "ActionClassificationModel", + "add_rotated_rects", "AnomalyDetection", "AnomalyResult", + "classification_models", "ClassificationModel", + "ClassificationResult", "Contour", + "detection_models", + "DetectedKeypoints", "DetectionModel", + "DetectionResult", + "get_contours", "ImageModel", "ImageResultWithSoftPrediction", "InstanceSegmentationResult", - "VisualPromptingResult", - "ZSLVisualPromptingResult", - "PredictedMask", - "SAMVisualPrompter", - "SAMLearnableVisualPrompter", "KeypointDetectionModel", - "TopDownKeypointDetectionPipeline", + "Label", "MaskRCNNModel", "Model", "OutputTransform", + "PredictedMask", + "Prompt", + "RotatedSegmentationResult", + "SAMDecoder", + "SAMImageEncoder", + "SAMLearnableVisualPrompter", + "SAMVisualPrompter", "SalientObjectDetectionModel", + "segmentation_models", "SegmentationModel", "SSD", + "TopDownKeypointDetectionPipeline", + "VisualPromptingResult", "YOLO", - "YoloV3ONNX", - "YoloV4", + "YOLOF", + "YOLOv3ONNX", + "YOLOv4", "YOLOv5", "YOLOv8", - "YOLOF", "YOLOX", - "SAMDecoder", - "SAMImageEncoder", - "ClassificationResult", - "Prompt", - "DetectionResult", - "DetectedKeypoints", - "classification_models", - "detection_models", - "segmentation_models", - "RotatedSegmentationResult", - "add_rotated_rects", - "get_contours", + "ZSLVisualPromptingResult", + "YoloV3ONNX", + "YoloV4", ] diff --git a/model_api/python/model_api/models/action_classification.py b/model_api/python/model_api/models/action_classification.py index efed98ad..8af041ab 100644 --- a/model_api/python/model_api/models/action_classification.py +++ b/model_api/python/model_api/models/action_classification.py @@ -10,10 +10,9 @@ import numpy as np from model_api.adapters.utils import RESIZE_TYPES, InputTransform -from model_api.models.result import Label +from model_api.models.result import ClassificationResult, Label from .model import Model -from .result import ClassificationResult from .types import BooleanValue, ListValue, NumericalValue, StringValue from .utils import load_labels diff --git a/model_api/python/model_api/models/result/__init__.py b/model_api/python/model_api/models/result/__init__.py index 6f10071b..c14deaf1 100644 --- a/model_api/python/model_api/models/result/__init__.py +++ b/model_api/python/model_api/models/result/__init__.py @@ -3,21 +3,13 @@ # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -from .types import ( - AnomalyResult, - ClassificationResult, - Contour, - DetectedKeypoints, - DetectionResult, - ImageResultWithSoftPrediction, - InstanceSegmentationResult, - Label, - PredictedMask, - Result, - RotatedSegmentationResult, - VisualPromptingResult, - ZSLVisualPromptingResult, -) +from .anomaly import AnomalyResult +from .base import Result +from .classification import ClassificationResult, Label +from .detection import DetectionResult +from .keypoint import DetectedKeypoints +from .segmentation import Contour, ImageResultWithSoftPrediction, InstanceSegmentationResult, RotatedSegmentationResult +from .visual_prompting import PredictedMask, VisualPromptingResult, ZSLVisualPromptingResult __all__ = [ "AnomalyResult", diff --git a/model_api/python/model_api/models/result/types/anomaly.py b/model_api/python/model_api/models/result/anomaly.py similarity index 100% rename from model_api/python/model_api/models/result/types/anomaly.py rename to model_api/python/model_api/models/result/anomaly.py diff --git a/model_api/python/model_api/models/result/types/base.py b/model_api/python/model_api/models/result/base.py similarity index 100% rename from model_api/python/model_api/models/result/types/base.py rename to model_api/python/model_api/models/result/base.py diff --git a/model_api/python/model_api/models/result/types/classification.py b/model_api/python/model_api/models/result/classification.py similarity index 100% rename from model_api/python/model_api/models/result/types/classification.py rename to model_api/python/model_api/models/result/classification.py diff --git a/model_api/python/model_api/models/result/types/detection.py b/model_api/python/model_api/models/result/detection.py similarity index 100% rename from model_api/python/model_api/models/result/types/detection.py rename to model_api/python/model_api/models/result/detection.py diff --git a/model_api/python/model_api/models/result/types/keypoint.py b/model_api/python/model_api/models/result/keypoint.py similarity index 100% rename from model_api/python/model_api/models/result/types/keypoint.py rename to model_api/python/model_api/models/result/keypoint.py diff --git a/model_api/python/model_api/models/result/types/segmentation.py b/model_api/python/model_api/models/result/segmentation.py similarity index 100% rename from model_api/python/model_api/models/result/types/segmentation.py rename to model_api/python/model_api/models/result/segmentation.py diff --git a/model_api/python/model_api/models/result/types/__init__.py b/model_api/python/model_api/models/result/types/__init__.py deleted file mode 100644 index 5b33d9fb..00000000 --- a/model_api/python/model_api/models/result/types/__init__.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Result types.""" - -# Copyright (C) 2024 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -from .anomaly import AnomalyResult -from .base import Result -from .classification import ClassificationResult, Label -from .detection import DetectionResult -from .keypoint import DetectedKeypoints -from .segmentation import ( - Contour, - ImageResultWithSoftPrediction, - InstanceSegmentationResult, - RotatedSegmentationResult, -) -from .visual_prompting import PredictedMask, VisualPromptingResult, ZSLVisualPromptingResult - -__all__ = [ - "AnomalyResult", - "ClassificationResult", - "Contour", - "DetectionResult", - "DetectedKeypoints", - "Label", - "ImageResultWithSoftPrediction", - "InstanceSegmentationResult", - "PredictedMask", - "Result", - "VisualPromptingResult", - "ZSLVisualPromptingResult", - "RotatedSegmentationResult", -] diff --git a/model_api/python/model_api/models/result/types/utils.py b/model_api/python/model_api/models/result/utils.py similarity index 100% rename from model_api/python/model_api/models/result/types/utils.py rename to model_api/python/model_api/models/result/utils.py diff --git a/model_api/python/model_api/models/result/types/visual_prompting.py b/model_api/python/model_api/models/result/visual_prompting.py similarity index 100% rename from model_api/python/model_api/models/result/types/visual_prompting.py rename to model_api/python/model_api/models/result/visual_prompting.py diff --git a/model_api/python/model_api/visualizer/__init__.py b/model_api/python/model_api/visualizer/__init__.py index ad194fef..ed8b255c 100644 --- a/model_api/python/model_api/visualizer/__init__.py +++ b/model_api/python/model_api/visualizer/__init__.py @@ -3,7 +3,9 @@ # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 +from .layout import Flatten, Layout +from .primitive import Overlay from .scene import Scene from .visualizer import Visualizer -__all__ = ["Scene", "Visualizer"] +__all__ = ["Overlay", "Scene", "Visualizer", "Layout", "Flatten"] diff --git a/model_api/python/model_api/visualizer/layout.py b/model_api/python/model_api/visualizer/layout.py index aa18490a..39e3117d 100644 --- a/model_api/python/model_api/visualizer/layout.py +++ b/model_api/python/model_api/visualizer/layout.py @@ -5,8 +5,8 @@ from __future__ import annotations -from abc import ABC -from typing import TYPE_CHECKING +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Type if TYPE_CHECKING: import PIL @@ -19,10 +19,33 @@ class Layout(ABC): """Base class for layouts.""" - def _compute_on_primitive(self, primitive: Primitive, image: PIL.Image, scene: Scene) -> PIL.Image | None: - if scene.has_primitives(type(primitive)): - primitives = scene.get_primitives(type(primitive)) - for primitive in primitives: - image = primitive.compute(image) + def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, scene: Scene) -> PIL.Image | None: + if scene.has_primitives(primitive): + primitives = scene.get_primitives(primitive) + for _primitive in primitives: + image = _primitive.compute(image) return image return None + + @abstractmethod + def __call__(self, scene: Scene) -> PIL.Image: + """Compute the layout.""" + + +class Flatten(Layout): + """Put all primitives on top of each other. + + Args: + *args (Type[Primitive]): Primitives to be applied. + """ + + def __init__(self, *args: Type[Primitive]) -> None: + self.children = args + + def __call__(self, scene: Scene) -> PIL.Image: + _image: PIL.Image = scene.base.copy() + for child in self.children: + output = self._compute_on_primitive(child, _image, scene) + if output is not None: + _image = output + return _image diff --git a/model_api/python/model_api/visualizer/primitive.py b/model_api/python/model_api/visualizer/primitive.py index 7afaee24..d24fc217 100644 --- a/model_api/python/model_api/visualizer/primitive.py +++ b/model_api/python/model_api/visualizer/primitive.py @@ -6,15 +6,38 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import TYPE_CHECKING -if TYPE_CHECKING: - import PIL +import numpy as np +import PIL class Primitive(ABC): """Primitive class.""" @abstractmethod - def compute(self, image: PIL.Image, **kwargs) -> PIL.Image: + def compute(self, image: PIL.Image) -> PIL.Image: pass + + +class Overlay(Primitive): + """Overlay primitive. + + Useful for XAI and Anomaly Maps. + + Args: + image (PIL.Image | np.ndarray): Image to be overlaid. + opacity (float): Opacity of the overlay. + """ + + def __init__(self, image: PIL.Image | np.ndarray, opacity: float = 0.4) -> None: + self.image = self._to_pil(image) + self.opacity = opacity + + def _to_pil(self, image: PIL.Image | np.ndarray) -> PIL.Image: + if isinstance(image, np.ndarray): + return PIL.Image.fromarray(image) + return image + + def compute(self, image: PIL.Image) -> PIL.Image: + _image = self.image.resize(image.size) + return PIL.Image.blend(image, _image, self.opacity) diff --git a/model_api/python/model_api/visualizer/scene.py b/model_api/python/model_api/visualizer/scene.py deleted file mode 100644 index e8065365..00000000 --- a/model_api/python/model_api/visualizer/scene.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Scene object.""" - -# Copyright (C) 2024 Intel Corporation -# SPDX-License-Identifier: Apache-2.0 - -from abc import abstractmethod -from typing import Type, Union - -import PIL - -from .layout import Layout -from .primitive import Primitive - - -class Scene: - """Scene object. - - Used by the visualizer to render. - """ - - def __init__( - self, - base: PIL.Image, - layout: Union[Layout, list[Layout], None] = None, - ) -> None: ... - - def show(self) -> PIL.Image: ... - - def save(self, path: str) -> None: ... - - def has_primitives(self, primitive: Type[Primitive]) -> bool: - return False - - def get_primitives(self, primitive: Type[Primitive]) -> list[Primitive]: - return [] - - @property - @abstractmethod - def default_layout(self) -> Layout: - """Default layout for the media.""" diff --git a/model_api/python/model_api/visualizer/scene/__init__.py b/model_api/python/model_api/visualizer/scene/__init__.py new file mode 100644 index 00000000..84469928 --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/__init__.py @@ -0,0 +1,22 @@ +"""Result visualization Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from .anomaly import AnomalyScene +from .classification import ClassificationScene +from .detection import DetectionScene +from .keypoint import KeypointScene +from .scene import Scene +from .segmentation import SegmentationScene +from .visual_prompting import VisualPromptingScene + +__all__ = [ + "AnomalyScene", + "ClassificationScene", + "DetectionScene", + "KeypointScene", + "Scene", + "SegmentationScene", + "VisualPromptingScene", +] diff --git a/model_api/python/model_api/visualizer/scene/anomaly.py b/model_api/python/model_api/visualizer/scene/anomaly.py new file mode 100644 index 00000000..7bb02756 --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/anomaly.py @@ -0,0 +1,24 @@ +"""Anomaly Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from PIL import Image + +from model_api.models.result import AnomalyResult +from model_api.visualizer.layout import Flatten, Layout +from model_api.visualizer.primitive import Overlay + +from .scene import Scene + + +class AnomalyScene(Scene): + """Anomaly Scene.""" + + def __init__(self, image: Image, result: AnomalyResult) -> None: + self.image = image + self.result = result + + @property + def default_layout(self) -> Layout: + return Flatten(Overlay) diff --git a/model_api/python/model_api/visualizer/scene/classification.py b/model_api/python/model_api/visualizer/scene/classification.py new file mode 100644 index 00000000..54f51ba7 --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/classification.py @@ -0,0 +1,24 @@ +"""Classification Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +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 .scene import Scene + + +class ClassificationScene(Scene): + """Classification Scene.""" + + def __init__(self, image: Image, result: ClassificationResult) -> None: + self.image = image + self.result = result + + @property + def default_layout(self) -> Layout: + return Flatten(Overlay) diff --git a/model_api/python/model_api/visualizer/scene/detection.py b/model_api/python/model_api/visualizer/scene/detection.py new file mode 100644 index 00000000..35e55bc2 --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/detection.py @@ -0,0 +1,18 @@ +"""Detection Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from PIL import Image + +from model_api.models.result import DetectionResult + +from .scene import Scene + + +class DetectionScene(Scene): + """Detection Scene.""" + + def __init__(self, image: Image, result: DetectionResult) -> None: + self.image = image + self.result = result diff --git a/model_api/python/model_api/visualizer/scene/keypoint.py b/model_api/python/model_api/visualizer/scene/keypoint.py new file mode 100644 index 00000000..3e34711c --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/keypoint.py @@ -0,0 +1,21 @@ +"""Keypoint Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from model_api.models.result import DetectedKeypoints +from model_api.visualizer.layout import Flatten, Layout +from model_api.visualizer.primitive import Overlay + +from .scene import Scene + + +class KeypointScene(Scene): + """Keypoint Scene.""" + + def __init__(self, result: DetectedKeypoints) -> None: + self.result = result + + @property + def default_layout(self) -> Layout: + return Flatten(Overlay) diff --git a/model_api/python/model_api/visualizer/scene/scene.py b/model_api/python/model_api/visualizer/scene/scene.py new file mode 100644 index 00000000..9fdb7072 --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/scene.py @@ -0,0 +1,72 @@ +"""Scene object.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from PIL import Image + +from model_api.visualizer.primitive import Overlay, Primitive + +if TYPE_CHECKING: + from pathlib import Path + + from model_api.visualizer.layout import Layout + + +class Scene: + """Scene object. + + Used by the visualizer to render. + """ + + def __init__( + self, + base: Image, + overlay: Overlay | list[Overlay] | np.ndarray | None = None, + layout: Layout | None = None, + ) -> None: + self.base = base + self.overlay = self._to_overlay(overlay) + self.layout = layout + + def show(self) -> Image: ... + + def save(self, path: Path) -> None: ... + + def render(self) -> Image: + if self.layout is None: + return self.default_layout(self) + return self.layout(self) + + def has_primitives(self, primitive: type[Primitive]) -> bool: + if primitive == Overlay: + return bool(self.overlay) + return False + + def get_primitives(self, primitive: type[Primitive]) -> list[Primitive]: + primitives: list[Primitive] | None = None + if primitive == Overlay: + primitives = self.overlay # type: ignore[assignment] # TODO(ashwinvaidya17): Address this in the next PR + if primitives is None: + msg = f"Primitive {primitive} not found" + raise ValueError(msg) + return primitives + + @property + def default_layout(self) -> Layout: + """Default layout for the media.""" + msg = "Default layout not implemented" + raise NotImplementedError(msg) + + def _to_overlay(self, overlay: Overlay | list[Overlay] | np.ndarray | None) -> list[Overlay] | None: + if isinstance(overlay, np.ndarray): + image = Image.fromarray(overlay) + return [Overlay(image)] + if isinstance(overlay, Overlay): + return [overlay] + return overlay diff --git a/model_api/python/model_api/visualizer/scene/segmentation.py b/model_api/python/model_api/visualizer/scene/segmentation.py new file mode 100644 index 00000000..e666804e --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/segmentation.py @@ -0,0 +1,15 @@ +"""Segmentation Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from model_api.models.result import InstanceSegmentationResult + +from .scene import Scene + + +class SegmentationScene(Scene): + """Segmentation Scene.""" + + def __init__(self, result: InstanceSegmentationResult) -> None: + self.result = result diff --git a/model_api/python/model_api/visualizer/scene/visual_prompting.py b/model_api/python/model_api/visualizer/scene/visual_prompting.py new file mode 100644 index 00000000..c31e0dd2 --- /dev/null +++ b/model_api/python/model_api/visualizer/scene/visual_prompting.py @@ -0,0 +1,15 @@ +"""Visual Prompting Scene.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +from model_api.models.result import VisualPromptingResult + +from .scene import Scene + + +class VisualPromptingScene(Scene): + """Visual Prompting Scene.""" + + def __init__(self, result: VisualPromptingResult) -> None: + self.result = result diff --git a/model_api/python/model_api/visualizer/visualizer.py b/model_api/python/model_api/visualizer/visualizer.py index 9519a9a1..8a2ea87e 100644 --- a/model_api/python/model_api/visualizer/visualizer.py +++ b/model_api/python/model_api/visualizer/visualizer.py @@ -3,10 +3,43 @@ # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 -import PIL +from pathlib import Path -from model_api.models.result import Result +from PIL import Image + +from model_api.models.result import ( + AnomalyResult, + ClassificationResult, + DetectionResult, + Result, +) + +from .layout import Layout +from .scene import AnomalyScene, ClassificationScene, DetectionScene, Scene class Visualizer: - def show(self, image: PIL.Image, result: Result) -> PIL.Image: ... + def __init__(self, layout: Layout) -> None: + self.layout = layout + + def show(self, image: Image, result: Result) -> Image: + scene = self._scene_from_result(image, result) + return scene.show() + + def save(self, image: Image, result: Result, path: Path) -> None: + scene = self._scene_from_result(image, result) + scene.save(path) + + def _scene_from_result(self, image: Image, result: Result) -> Scene: + scene: Scene + if isinstance(result, AnomalyResult): + scene = AnomalyScene(image, result) + elif isinstance(result, ClassificationResult): + scene = ClassificationScene(image, result) + elif isinstance(result, DetectionResult): + scene = DetectionScene(image, result) + else: + msg = f"Unsupported result type: {type(result)}" + raise ValueError(msg) + + return scene diff --git a/model_api/python/model_api/models/result/scene/__init__.py b/tests/python/unit/visualizer/__init__.py similarity index 69% rename from model_api/python/model_api/models/result/scene/__init__.py rename to tests/python/unit/visualizer/__init__.py index bf383d1a..8fedea4d 100644 --- a/model_api/python/model_api/models/result/scene/__init__.py +++ b/tests/python/unit/visualizer/__init__.py @@ -1,4 +1,4 @@ -"""Result visualization Scene.""" +"""Visualization tests.""" # Copyright (C) 2024 Intel Corporation # SPDX-License-Identifier: Apache-2.0 diff --git a/tests/python/unit/visualizer/conftest.py b/tests/python/unit/visualizer/conftest.py new file mode 100644 index 00000000..7e04d138 --- /dev/null +++ b/tests/python/unit/visualizer/conftest.py @@ -0,0 +1,28 @@ +"""Conftest for visualization tests.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import pytest +from PIL import Image + +from model_api.visualizer import Overlay, Scene + + +@pytest.fixture(scope="session") +def mock_image(): + data = np.zeros((100, 100, 3), dtype=np.uint8) + data *= 255 + return Image.fromarray(data) + + +@pytest.fixture(scope="session") +def mock_scene(mock_image: Image) -> Scene: + """Mock scene.""" + overlay = np.zeros((100, 100, 3), dtype=np.uint8) + overlay[50, 50] = [255, 0, 0] + return Scene( + base=mock_image, + overlay=Overlay(overlay), + ) diff --git a/tests/python/unit/visualizer/test_layout.py b/tests/python/unit/visualizer/test_layout.py new file mode 100644 index 00000000..b4502f3b --- /dev/null +++ b/tests/python/unit/visualizer/test_layout.py @@ -0,0 +1,27 @@ +"""Test layout.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +from PIL import Image + +from model_api.visualizer import Flatten, Scene +from model_api.visualizer.primitive import Overlay + + +def test_flatten_layout(mock_image: Image, mock_scene: Scene): + """Test if the layout is created correctly.""" + overlay = np.zeros((100, 100, 3), dtype=np.uint8) + overlay[50, 50] = [255, 0, 0] + overlay = Image.fromarray(overlay) + + expected_image = Image.blend(mock_image, overlay, 0.4) + mock_scene.layout = Flatten(Overlay) + assert mock_scene.render() == expected_image + + +def test_flatten_layout_with_no_primitives(mock_image: Image, mock_scene: Scene): + """Test if the layout is created correctly.""" + mock_scene.layout = Flatten() + assert mock_scene.render() == mock_image diff --git a/tests/python/unit/visualizer/test_primitive.py b/tests/python/unit/visualizer/test_primitive.py new file mode 100644 index 00000000..9f6a210b --- /dev/null +++ b/tests/python/unit/visualizer/test_primitive.py @@ -0,0 +1,24 @@ +"""Tests for primitives.""" + +# Copyright (C) 2024 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import PIL + +from model_api.visualizer import Overlay + + +def test_overlay(mock_image: PIL.Image): + """Test if the overlay is created correctly.""" + empty_image = PIL.Image.new("RGB", (100, 100)) + expected_image = PIL.Image.blend(empty_image, mock_image, 0.4) + # Test from image + overlay = Overlay(mock_image) + assert overlay.compute(empty_image) == expected_image + + # Test from numpy array + data = np.zeros((100, 100, 3), dtype=np.uint8) + data *= 255 + overlay = Overlay(data) + assert overlay.compute(empty_image) == expected_image