115 lines
3.9 KiB
Python
115 lines
3.9 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
from app.config import settings
|
|
from app.pipeline import BirdCascadePipeline
|
|
|
|
|
|
class PipelineSwitchTests(unittest.TestCase):
|
|
setting_names = (
|
|
"enable_detection",
|
|
"enable_classification",
|
|
"enable_density",
|
|
"uniform_species_color",
|
|
)
|
|
|
|
def setUp(self):
|
|
self.original_values = {
|
|
name: getattr(settings, name)
|
|
for name in self.setting_names
|
|
}
|
|
|
|
def tearDown(self):
|
|
for name, value in self.original_values.items():
|
|
object.__setattr__(settings, name, value)
|
|
|
|
def test_all_models_can_be_disabled_without_loading_onnx(self):
|
|
object.__setattr__(settings, "enable_detection", False)
|
|
object.__setattr__(settings, "enable_classification", False)
|
|
object.__setattr__(settings, "enable_density", False)
|
|
|
|
with (
|
|
patch("app.pipeline.YOLOv5ONNXDetector") as detector_mock,
|
|
patch("app.pipeline.ConvNeXtONNXClassifier") as classifier_mock,
|
|
patch("app.pipeline.DensityEstimatorONNX") as density_mock,
|
|
):
|
|
pipeline = BirdCascadePipeline()
|
|
|
|
self.assertIsNone(pipeline.detector)
|
|
self.assertIsNone(pipeline.classifier)
|
|
self.assertIsNone(pipeline.density_estimator)
|
|
detector_mock.assert_not_called()
|
|
classifier_mock.assert_not_called()
|
|
density_mock.assert_not_called()
|
|
|
|
output = pipeline.predict(
|
|
np.zeros((32, 32, 3), dtype=np.uint8),
|
|
)
|
|
self.assertEqual(output["results"], [])
|
|
self.assertIsNone(output["density"])
|
|
|
|
def test_detection_without_classification_keeps_detection_fields(self):
|
|
object.__setattr__(settings, "enable_detection", True)
|
|
object.__setattr__(settings, "enable_classification", False)
|
|
object.__setattr__(settings, "enable_density", False)
|
|
|
|
fake_detector_result = [
|
|
{
|
|
"bbox_xyxy": [4.0, 4.0, 20.0, 20.0],
|
|
"det_confidence": 0.9,
|
|
"det_class_id": 0,
|
|
"det_class_name": "bird",
|
|
}
|
|
]
|
|
|
|
with (
|
|
patch(
|
|
"app.pipeline.read_classes",
|
|
return_value=["bird"],
|
|
),
|
|
patch(
|
|
"app.pipeline.os.path.exists",
|
|
return_value=True,
|
|
),
|
|
patch("app.pipeline.YOLOv5ONNXDetector") as detector_class,
|
|
patch("app.pipeline.ConvNeXtONNXClassifier") as classifier_class,
|
|
):
|
|
detector_class.return_value.infer.return_value = fake_detector_result
|
|
pipeline = BirdCascadePipeline()
|
|
output = pipeline.predict(
|
|
np.zeros((32, 32, 3), dtype=np.uint8),
|
|
)
|
|
|
|
classifier_class.assert_not_called()
|
|
self.assertEqual(len(output["results"]), 1)
|
|
result = output["results"][0]
|
|
self.assertEqual(result["det_class_name"], "bird")
|
|
self.assertEqual(result["cls_class_id"], -1)
|
|
self.assertEqual(result["cls_class_name"], "")
|
|
self.assertEqual(result["cls_confidence"], 0.0)
|
|
self.assertFalse(result["show_cls_name"])
|
|
|
|
def test_uniform_species_color_uses_one_color_for_every_class(self):
|
|
pipeline = object.__new__(BirdCascadePipeline)
|
|
|
|
object.__setattr__(settings, "uniform_species_color", False)
|
|
first_color = pipeline._get_color_by_class_id(1)
|
|
second_color = pipeline._get_color_by_class_id(2)
|
|
self.assertNotEqual(first_color, second_color)
|
|
|
|
object.__setattr__(settings, "uniform_species_color", True)
|
|
self.assertEqual(
|
|
pipeline._get_color_by_class_id(1),
|
|
pipeline._get_color_by_class_id(2),
|
|
)
|
|
self.assertEqual(
|
|
pipeline._get_color_by_class_id(999),
|
|
(0, 190, 0),
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|