Skip to content

Commit

Permalink
Enhance Flatten layout handling: update image computation to handle N…
Browse files Browse the repository at this point in the history
…one outputs and add unit test for empty layout scenario.
  • Loading branch information
ashwinvaidya17 committed Jan 3, 2025
1 parent 27669f5 commit 6d158bd
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 1 deletion.
4 changes: 3 additions & 1 deletion model_api/python/model_api/visualizer/layout.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,5 +45,7 @@ def __init__(self, *args: Type[Primitive]) -> None:
def __call__(self, scene: Scene) -> PIL.Image:
_image: PIL.Image = scene.base.copy()
for child in self.children:
_image = self._compute_on_primitive(child, _image, scene)
output = self._compute_on_primitive(child, _image, scene)
if output is not None:
_image = output
return _image
6 changes: 6 additions & 0 deletions tests/python/unit/visualizer/test_layout.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,9 @@ def test_flatten_layout(mock_image: Image, mock_scene: Scene):
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

0 comments on commit 6d158bd

Please sign in to comment.