Skip to content

Commit

Permalink
Add tests
Browse files Browse the repository at this point in the history
Signed-off-by: Ashwin Vaidya <[email protected]>
  • Loading branch information
ashwinvaidya17 committed Jan 22, 2025
1 parent 56feddf commit 4a5ac67
Showing 1 changed file with 15 additions and 1 deletion.
16 changes: 15 additions & 1 deletion tests/python/unit/visualizer/test_scene.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()

0 comments on commit 4a5ac67

Please sign in to comment.