Spaces:
Running
Running
File size: 1,211 Bytes
67f4974 b41850c 67f4974 b41850c 67f4974 b41850c 67f4974 b41850c 67f4974 b41850c 67f4974 b41850c 67f4974 b41850c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
"""
Testing Detection module
"""
import os
import pytest
from detection import ml_detection
# Test model loading
@pytest.mark.parametrize(
"test_model_uri",
[
("facebook/detr-resnet-50"),
("facebook/detr-resnet-101"),
],
)
def test_load_model(test_model_uri):
"""Testing model loading"""
processor, model = ml_detection.load_model(test_model_uri)
assert processor is not None
assert model is not None
# Test image detection
@pytest.mark.parametrize(
"test_model_uri",
[
("facebook/detr-resnet-50"),
("facebook/detr-resnet-101"),
],
)
def test_object_detection(test_model_uri):
"""Testing object detection function"""
processor, model = ml_detection.load_model(test_model_uri)
# Get the directory of the current test file
test_dir = os.path.dirname(os.path.abspath(__file__))
# Construct the image path relative to the test directory
image_path = os.path.join(test_dir, "data", "savanna.jpg")
with open(image_path, "rb") as f:
image_bytes = f.read()
results = ml_detection.object_detection(processor, model, image_bytes)
assert results is not None
assert isinstance(results, dict)
|