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()