From 6d158bdffcf1691ed3aa0fb6114c2e62a19003ec Mon Sep 17 00:00:00 2001 From: Ashwin Vaidya Date: Fri, 3 Jan 2025 13:21:16 +0100 Subject: [PATCH] Enhance Flatten layout handling: update image computation to handle None outputs and add unit test for empty layout scenario. --- model_api/python/model_api/visualizer/layout.py | 4 +++- tests/python/unit/visualizer/test_layout.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/model_api/python/model_api/visualizer/layout.py b/model_api/python/model_api/visualizer/layout.py index e4909507..39e3117d 100644 --- a/model_api/python/model_api/visualizer/layout.py +++ b/model_api/python/model_api/visualizer/layout.py @@ -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 diff --git a/tests/python/unit/visualizer/test_layout.py b/tests/python/unit/visualizer/test_layout.py index 810fd5fd..b4502f3b 100644 --- a/tests/python/unit/visualizer/test_layout.py +++ b/tests/python/unit/visualizer/test_layout.py @@ -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