mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-09-26 15:09:00 +03:00
95 lines
3.1 KiB
Python
95 lines
3.1 KiB
Python
"""Tests for parsing detection hardware device strings."""
|
|
|
|
import unittest
|
|
|
|
from frigate.detectors.detector_config import ModelConfig
|
|
from frigate.detectors.device import (
|
|
DeviceParseError,
|
|
build_detector_config,
|
|
parse_device,
|
|
runner_names,
|
|
)
|
|
|
|
|
|
class TestParseDevice(unittest.TestCase):
|
|
def test_bare_detector_has_no_device(self):
|
|
spec = parse_device("cpu")
|
|
|
|
self.assertEqual(spec.detector, "cpu")
|
|
self.assertIsNone(spec.device)
|
|
|
|
def test_device_is_everything_after_the_first_colon(self):
|
|
spec = parse_device("edgetpu:pci:0")
|
|
|
|
self.assertEqual(spec.detector, "edgetpu")
|
|
self.assertEqual(spec.device, "pci:0")
|
|
|
|
def test_trailing_colon_keeps_an_empty_device(self):
|
|
# an empty edgetpu device selects a native Coral
|
|
spec = parse_device("edgetpu:")
|
|
|
|
self.assertEqual(spec.detector, "edgetpu")
|
|
self.assertEqual(spec.device, "")
|
|
|
|
def test_unknown_detector_is_rejected(self):
|
|
with self.assertRaises(DeviceParseError):
|
|
parse_device("not_a_detector:0")
|
|
|
|
def test_device_that_the_detector_cannot_use_is_rejected(self):
|
|
# tensorrt takes a gpu index
|
|
with self.assertRaises(DeviceParseError):
|
|
parse_device("tensorrt:the-fast-one")
|
|
|
|
|
|
class TestBuildDetectorConfig(unittest.TestCase):
|
|
def _build(self, raw: str):
|
|
return build_detector_config(parse_device(raw), ModelConfig())
|
|
|
|
def test_device_lands_on_the_detector_field(self):
|
|
for raw, expected in (
|
|
("edgetpu:usb", "usb"),
|
|
("edgetpu:pci:1", "pci:1"),
|
|
("openvino:GPU.1", "GPU.1"),
|
|
("onnx:CPU", "CPU"),
|
|
("memryx:PCIe:0", "PCIe:0"),
|
|
):
|
|
with self.subTest(raw=raw):
|
|
self.assertEqual(self._build(raw).device, expected)
|
|
|
|
def test_detectors_that_name_the_field_something_else(self):
|
|
self.assertEqual(self._build("cpu:4").num_threads, 4)
|
|
self.assertEqual(self._build("rknn:2").num_cores, 2)
|
|
|
|
def test_device_is_coerced_to_the_detector_field_type(self):
|
|
self.assertEqual(self._build("tensorrt:1").device, 1)
|
|
|
|
def test_omitted_device_falls_back_to_the_detector_default(self):
|
|
self.assertEqual(self._build("cpu").num_threads, 3)
|
|
self.assertEqual(self._build("rknn").num_cores, 0)
|
|
self.assertEqual(self._build("openvino").device, "AUTO")
|
|
self.assertIsNone(self._build("edgetpu").device)
|
|
|
|
def test_the_model_is_attached(self):
|
|
model = ModelConfig(path="/cpu_model.tflite")
|
|
|
|
self.assertIs(build_detector_config(parse_device("cpu"), model).model, model)
|
|
|
|
|
|
class TestRunnerNames(unittest.TestCase):
|
|
def test_unique_devices_keep_their_name(self):
|
|
devices = [parse_device("edgetpu:pci:0"), parse_device("edgetpu:pci:1")]
|
|
|
|
self.assertEqual(runner_names(devices), ["edgetpu:pci:0", "edgetpu:pci:1"])
|
|
|
|
def test_repeated_devices_are_numbered(self):
|
|
devices = [parse_device("openvino:GPU")] * 3
|
|
|
|
self.assertEqual(
|
|
runner_names(devices),
|
|
["openvino:GPU", "openvino:GPU#2", "openvino:GPU#3"],
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|