diff --git a/frigate/test/test_object_detector.py b/frigate/test/test_object_detector.py new file mode 100644 index 000000000..0a3137273 --- /dev/null +++ b/frigate/test/test_object_detector.py @@ -0,0 +1,94 @@ +import unittest +from unittest.mock import patch + +import numpy as np +from frigate.config import DetectorTypeEnum +import frigate.object_detection + + +class TestLocalObjectDetector(unittest.TestCase): + @patch("frigate.object_detection.EdgeTpuTfl") + @patch("frigate.object_detection.CpuTfl") + def test_localdetectorprocess_given_type_cpu_should_call_cputfl_init( + self, mock_cputfl, mock_edgetputfl + ): + test_obj = frigate.object_detection.LocalObjectDetector( + det_type=DetectorTypeEnum.cpu, model_path="/test/modelpath", num_threads=6 + ) + + assert test_obj is not None + mock_edgetputfl.assert_not_called() + mock_cputfl.assert_called_once_with(model_path="/test/modelpath", num_threads=6) + + @patch("frigate.object_detection.EdgeTpuTfl") + @patch("frigate.object_detection.CpuTfl") + def test_localdetectorprocess_given_type_edgtpu_should_call_edgtpu_init( + self, mock_cputfl, mock_edgetputfl + ): + test_obj = frigate.object_detection.LocalObjectDetector( + det_type=DetectorTypeEnum.edgetpu, + det_device="usb", + model_path="/test/modelpath", + ) + + assert test_obj is not None + mock_cputfl.assert_not_called() + mock_edgetputfl.assert_called_once_with( + det_device="usb", model_path="/test/modelpath" + ) + + @patch("frigate.object_detection.CpuTfl") + def test_detect_raw_given_tensor_input_should_return_api_detect_raw_result( + self, mock_cputfl + ): + TEST_DATA = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + TEST_DETECT_RESULT = np.ndarray([1, 2, 4, 8, 16, 32]) + test_obj_detect = frigate.object_detection.LocalObjectDetector( + det_device=DetectorTypeEnum.cpu + ) + + mock_det_api = mock_cputfl.return_value + mock_det_api.detect_raw.return_value = TEST_DETECT_RESULT + + test_result = test_obj_detect.detect_raw(TEST_DATA) + + mock_det_api.detect_raw.assert_called_once_with(tensor_input=TEST_DATA) + assert test_result is mock_det_api.detect_raw.return_value + + @patch("frigate.object_detection.CpuTfl") + @patch("frigate.object_detection.load_labels") + def test_detect_given_tensor_input_should_return_lfiltered_detections( + self, mock_load_labels, mock_cputfl + ): + TEST_DATA = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + TEST_DETECT_RAW = [ + [2, 0.9, 5, 4, 3, 2], + [1, 0.5, 8, 7, 6, 5], + [0, 0.4, 2, 4, 8, 16], + ] + TEST_DETECT_RESULT = [ + ("label-3", 0.9, (5, 4, 3, 2)), + ("label-2", 0.5, (8, 7, 6, 5)), + ] + TEST_LABEL_FILE = "/test_labels.txt" + mock_load_labels.return_value = [ + "label-1", + "label-2", + "label-3", + "label-4", + "label-5", + ] + + test_obj_detect = frigate.object_detection.LocalObjectDetector( + det_device=DetectorTypeEnum.cpu, labels=TEST_LABEL_FILE + ) + + mock_load_labels.assert_called_once_with(TEST_LABEL_FILE) + + mock_det_api = mock_cputfl.return_value + mock_det_api.detect_raw.return_value = TEST_DETECT_RAW + + test_result = test_obj_detect.detect(tensor_input=TEST_DATA, threshold=0.5) + + mock_det_api.detect_raw.assert_called_once_with(tensor_input=TEST_DATA) + assert test_result == TEST_DETECT_RESULT