mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-30 19:09:13 +01:00
4383b883c0
* Refactor EdgeTPU and CPU model handling to detector submodules. * Fix selecting the correct detection device type from the config * Remove detector type check when creating ObjectDetectProcess * Fixes after rebasing to 0.11 * Add init file to detector folder * Rename to detect_api Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com> * Add unit test for LocalObjectDetector class * Add configuration for model inputs Support transforming detection regions to RGB or BGR. Support specifying the input tensor shape. The tensor shape has a standard format ["BHWC"] when handed to the detector, but can be transformed in the detector to match the model shape using the model input_tensor config. * Add documentation for new model config parameters * Add input tensor transpose to LocalObjectDetector * Change the model input tensor config to use an enumeration * Updates for model config documentation Co-authored-by: Nicolas Mowen <nickmowen213@gmail.com>
131 lines
4.6 KiB
Python
131 lines
4.6 KiB
Python
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
from frigate.config import DetectorTypeEnum, InputTensorEnum, ModelConfig
|
|
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_cfg = ModelConfig()
|
|
test_cfg.path = "/test/modelpath"
|
|
test_obj = frigate.object_detection.LocalObjectDetector(
|
|
det_type=DetectorTypeEnum.cpu, model_config=test_cfg, num_threads=6
|
|
)
|
|
|
|
assert test_obj is not None
|
|
mock_edgetputfl.assert_not_called()
|
|
mock_cputfl.assert_called_once_with(model_config=test_cfg, 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_cfg = ModelConfig()
|
|
test_cfg.path = "/test/modelpath"
|
|
|
|
test_obj = frigate.object_detection.LocalObjectDetector(
|
|
det_type=DetectorTypeEnum.edgetpu,
|
|
det_device="usb",
|
|
model_config=test_cfg,
|
|
)
|
|
|
|
assert test_obj is not None
|
|
mock_cputfl.assert_not_called()
|
|
mock_edgetputfl.assert_called_once_with(det_device="usb", model_config=test_cfg)
|
|
|
|
@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")
|
|
def test_detect_raw_given_tensor_input_should_call_api_detect_raw_with_transposed_tensor(
|
|
self, mock_cputfl
|
|
):
|
|
TEST_DATA = np.zeros((1, 32, 32, 3), np.uint8)
|
|
TEST_DETECT_RESULT = np.ndarray([1, 2, 4, 8, 16, 32])
|
|
|
|
test_cfg = ModelConfig()
|
|
test_cfg.input_tensor = InputTensorEnum.nchw
|
|
|
|
test_obj_detect = frigate.object_detection.LocalObjectDetector(
|
|
det_device=DetectorTypeEnum.cpu, model_config=test_cfg
|
|
)
|
|
|
|
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()
|
|
assert (
|
|
mock_det_api.detect_raw.call_args.kwargs["tensor_input"].shape
|
|
== np.zeros((1, 3, 32, 32)).shape
|
|
)
|
|
|
|
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 = np.zeros((1, 32, 32, 3), np.uint8)
|
|
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,
|
|
model_config=ModelConfig(),
|
|
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()
|
|
assert (
|
|
mock_det_api.detect_raw.call_args.kwargs["tensor_input"].shape
|
|
== np.zeros((1, 32, 32, 3)).shape
|
|
)
|
|
assert test_result == TEST_DETECT_RESULT
|