Add HStack Layout (#250)

* Add HStack Layout

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

* Update copyright year in visualization modules to reflect 2024-2025

* Fix color assignment in HStack layout test: swap blue and red overlay colors

---------

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

Author: ashwinvaidya17

model_api/python/model_api/visualizer/__init__.py

@@ -3,9 +3,9 @@
 # Copyright (C) 2024 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
-from .layout import Flatten, Layout
+from .layout import Flatten, HStack, Layout
 from .primitive import Overlay
 from .scene import Scene
 from .visualizer import Visualizer
 
-__all__ = ["Overlay", "Scene", "Visualizer", "Layout", "Flatten"]
+__all__ = ["Overlay", "Scene", "Visualizer", "Layout", "Flatten", "HStack"]

model_api/python/model_api/visualizer/layout/__init__.py

@@ -0,0 +1,10 @@
+"""Visualization Layout."""
+
+# Copyright (C) 2024-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from .flatten import Flatten
+from .hstack import HStack
+from .layout import Layout
+
+__all__ = ["Flatten", "HStack", "Layout"]

model_api/python/model_api/visualizer/layout.py renamed to model_api/python/model_api/visualizer/layout/flatten.py

@@ -1,23 +1,30 @@
-"""Visualization Layout"""
+"""Flatten Layout."""
 
-# Copyright (C) 2024 Intel Corporation
+# Copyright (C) 2024-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Type
+from typing import TYPE_CHECKING, Type, Union
+
+from .layout import Layout
 
 if TYPE_CHECKING:
     import PIL
 
     from model_api.visualizer.primitive import Primitive
+    from model_api.visualizer.scene import Scene
 
-    from .scene import Scene
 
+class Flatten(Layout):
+    """Put all primitives on top of each other.
 
-class Layout(ABC):
-    """Base class for layouts."""
+    Args:
+        *args (Union[Type[Primitive], Layout]): Primitives or layouts to be applied.
+    """
+
+    def __init__(self, *args: Union[Type[Primitive], Layout]) -> None:
+        self.children = args
 
     def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, scene: Scene) -> PIL.Image | None:
         if scene.has_primitives(primitive):
@@ -27,25 +34,8 @@ def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, sc
             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
+            _image = child(scene) if isinstance(child, Layout) else self._compute_on_primitive(child, _image, scene)
         return _image

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

@@ -0,0 +1,78 @@
+"""Horizontal Stack Layout."""
+
+# Copyright (C) 2024-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Type, Union
+
+import PIL
+
+from .layout import Layout
+
+if TYPE_CHECKING:
+    from model_api.visualizer.primitive import Primitive
+    from model_api.visualizer.scene import Scene
+
+
+class HStack(Layout):
+    """Horizontal Stack Layout.
+
+    Args:
+        *args (Union[Type[Primitive], Layout]): Primitives or layouts to be applied.
+    """
+
+    def __init__(self, *args: Union[Type[Primitive], Layout]) -> None:
+        self.children = args
+
+    def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, scene: Scene) -> PIL.Image | None:
+        if scene.has_primitives(primitive):
+            images = []
+            for _primitive in scene.get_primitives(primitive):
+                _image = _primitive.compute(image.copy())
+                images.append(_image)
+            return self._stitch(*images)
+        return None
+
+    @staticmethod
+    def _stitch(*images: PIL.Image) -> PIL.Image:
+        """Stitch images horizontally.
+
+        Args:
+            images (PIL.Image): Images to be stitched.
+
+        Returns:
+            PIL.Image: Stitched image.
+        """
+        new_image = PIL.Image.new(
+            "RGB",
+            (
+                sum(image.width for image in images),
+                max(image.height for image in images),
+            ),
+        )
+        x_offset = 0
+        for image in images:
+            new_image.paste(image, (x_offset, 0))
+            x_offset += image.width
+        return new_image
+
+    def __call__(self, scene: Scene) -> PIL.Image:
+        """Stitch images horizontally.
+
+        Args:
+            scene (Scene): Scene to be stitched.
+
+        Returns:
+            PIL.Image: Stitched image.
+        """
+        images: list[PIL.Image] = []
+        for child in self.children:
+            if isinstance(child, Layout):
+                _image = child(scene)
+            else:
+                _image = self._compute_on_primitive(child, scene.base.copy(), scene)
+            if _image is not None:
+                images.append(_image)
+        return self._stitch(*images)

model_api/python/model_api/visualizer/layout/layout.py

@@ -0,0 +1,28 @@
+"""Visualization Layout."""
+
+# Copyright (C) 2024-2025 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Type
+
+if TYPE_CHECKING:
+    import PIL
+
+    from model_api.visualizer.primitive import Primitive
+
+    from .scene import Scene
+
+
+class Layout(ABC):
+    """Base class for layouts."""
+
+    @abstractmethod
+    def _compute_on_primitive(self, primitive: Type[Primitive], image: PIL.Image, scene: Scene) -> PIL.Image | None:
+        pass
+
+    @abstractmethod
+    def __call__(self, scene: Scene) -> PIL.Image:
+        """Compute the layout."""

tests/python/unit/visualizer/__init__.py

@@ -1,4 +1,4 @@
 """Visualization tests."""
 
-# Copyright (C) 2024 Intel Corporation
+# Copyright (C) 2024-2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0

tests/python/unit/visualizer/test_layout.py

@@ -6,7 +6,7 @@
 import numpy as np
 from PIL import Image
 
-from model_api.visualizer import Flatten, Scene
+from model_api.visualizer import Flatten, HStack, Scene
 from model_api.visualizer.primitive import Overlay
 
 
@@ -25,3 +25,26 @@ 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
+
+
+def test_hstack_layout():
+    """Test if the layout is created correctly."""
+    blue_overlay = np.zeros((100, 100, 3), dtype=np.uint8)
+    blue_overlay[50, 50] = [0, 0, 255]
+    blue_overlay = Image.fromarray(blue_overlay)
+
+    red_overlay = np.zeros((100, 100, 3), dtype=np.uint8)
+    red_overlay[50, 50] = [255, 0, 0]
+    red_overlay = Image.fromarray(red_overlay)
+
+    mock_scene = Scene(
+        base=Image.new("RGB", (100, 100)),
+        overlay=[Overlay(blue_overlay, opacity=1.0), Overlay(red_overlay, opacity=1.0)],
+        layout=HStack(Overlay),
+    )
+
+    expected_image = Image.new("RGB", (200, 100))
+    expected_image.paste(blue_overlay, (0, 0))
+    expected_image.paste(red_overlay, (100, 0))
+
+    assert mock_scene.render() == expected_image