From b7fd53ecc7989c773109795fee19665e5138b354 Mon Sep 17 00:00:00 2001 From: Dennis George Date: Sat, 10 Dec 2022 07:50:49 -0600 Subject: [PATCH] simplify config --- frigate/config.py | 4 ++-- frigate/detectors/config.py | 34 ++++------------------------ frigate/test/test_config.py | 5 ++-- frigate/test/test_object_detector.py | 18 ++++----------- 4 files changed, 15 insertions(+), 46 deletions(-) diff --git a/frigate/config.py b/frigate/config.py index 77b7253eb..d3a2d2983 100644 --- a/frigate/config.py +++ b/frigate/config.py @@ -9,7 +9,7 @@ from typing import Dict, List, Optional, Tuple, Union import matplotlib.pyplot as plt import numpy as np import yaml -from pydantic import BaseModel, Extra, Field, validator, parse_obj_as +from pydantic import BaseModel, Extra, Field, validator from pydantic.fields import PrivateAttr from frigate.const import ( @@ -972,7 +972,7 @@ class FrigateConfig(FrigateBaseModel): config.cameras[name] = camera_config for key, detector in config.detectors.items(): - detector_config: DetectorConfig = parse_obj_as(DetectorConfig, detector) + detector_config: DetectorConfig = DetectorConfig.parse_obj(detector) if detector_config.model is None: detector_config.model = config.model else: diff --git a/frigate/detectors/config.py b/frigate/detectors/config.py index 30c404211..3f666564b 100644 --- a/frigate/detectors/config.py +++ b/frigate/detectors/config.py @@ -57,41 +57,17 @@ class ModelConfig(BaseModel): extra = Extra.forbid -class BaseDetectorConfig(BaseModel): - type: DetectorTypeEnum = Field(default=DetectorTypeEnum.cpu, title="Detector Type") +class DetectorConfig(BaseModel): + type: str = Field(default=DetectorTypeEnum.cpu, title="Detector Type") model: ModelConfig = Field( default=None, title="Detector specific model configuration." ) + num_threads: Optional[int] = Field(default=3, title="Number of detection threads") + device: Optional[str] = Field(default="usb", title="Device Type") class Config: - extra = Extra.forbid + extra = Extra.allow arbitrary_types_allowed = True -class CpuDetectorConfig(BaseDetectorConfig): - type: Literal[DetectorTypeEnum.cpu] = Field( - default=DetectorTypeEnum.cpu, title="Detector Type" - ) - num_threads: int = Field(default=3, title="Number of detection threads") - - -class EdgeTpuDetectorConfig(BaseDetectorConfig): - type: Literal[DetectorTypeEnum.edgetpu] = Field( - default=DetectorTypeEnum.edgetpu, title="Detector Type" - ) - device: str = Field(default="usb", title="Device Type") - - -class OpenVinoDetectorConfig(BaseDetectorConfig): - type: Literal[DetectorTypeEnum.openvino] = Field( - default=DetectorTypeEnum.openvino, title="Detector Type" - ) - device: str = Field(default="usb", title="Device Type") - - -DetectorConfig = Annotated[ - Union[CpuDetectorConfig, EdgeTpuDetectorConfig, OpenVinoDetectorConfig], - Field(discriminator="type"), -] - DEFAULT_DETECTORS = parse_obj_as(Dict[str, DetectorConfig], {"cpu": {"type": "cpu"}}) diff --git a/frigate/test/test_config.py b/frigate/test/test_config.py index 9fa557985..61dd73847 100644 --- a/frigate/test/test_config.py +++ b/frigate/test/test_config.py @@ -44,12 +44,10 @@ class TestConfig(unittest.TestCase): "detectors": { "cpu": { "type": "cpu", - "num_threads": 5, "model": {"path": "/cpu_model.tflite"}, }, "edgetpu": { "type": "edgetpu", - "device": "usb", "model": {"path": "/edgetpu_model.tflite", "width": 160}, }, "openvino": { @@ -71,6 +69,9 @@ class TestConfig(unittest.TestCase): assert runtime_config.detectors["edgetpu"].type == DetectorTypeEnum.edgetpu assert runtime_config.detectors["openvino"].type == DetectorTypeEnum.openvino + assert runtime_config.detectors["cpu"].num_threads == 3 + assert runtime_config.detectors["edgetpu"].device == "usb" + assert runtime_config.model.path == "/default.tflite" assert runtime_config.detectors["cpu"].model.path == "/cpu_model.tflite" assert runtime_config.detectors["edgetpu"].model.path == "/edgetpu_model.tflite" diff --git a/frigate/test/test_object_detector.py b/frigate/test/test_object_detector.py index 586ee0979..6c1bdb295 100644 --- a/frigate/test/test_object_detector.py +++ b/frigate/test/test_object_detector.py @@ -2,16 +2,10 @@ import unittest from unittest.mock import Mock, patch import numpy as np -from pydantic import parse_obj_as from frigate.enums import InputTensorEnum from frigate.detectors import DetectorTypeEnum -from frigate.detectors.config import ( - CpuDetectorConfig, - EdgeTpuDetectorConfig, - DetectorConfig, - ModelConfig, -) +from frigate.detectors.config import DetectorConfig, ModelConfig import frigate.detectors as detectors import frigate.object_detection @@ -24,9 +18,7 @@ class TestLocalObjectDetector(unittest.TestCase): "frigate.detectors.api_types", {det_type: Mock() for det_type in DetectorTypeEnum}, ): - test_cfg = parse_obj_as( - DetectorConfig, {"type": det_type, "model": {}} - ) + test_cfg = DetectorConfig.parse_obj({"type": det_type, "model": {}}) test_cfg.model.path = "/test/modelpath" test_obj = frigate.object_detection.LocalObjectDetector( detector_config=test_cfg @@ -49,7 +41,7 @@ class TestLocalObjectDetector(unittest.TestCase): 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( - detector_config=CpuDetectorConfig(model=ModelConfig()) + detector_config=DetectorConfig(type="cpu", model=ModelConfig()) ) mock_det_api = mock_cputfl.return_value @@ -72,7 +64,7 @@ class TestLocalObjectDetector(unittest.TestCase): TEST_DATA = np.zeros((1, 32, 32, 3), np.uint8) TEST_DETECT_RESULT = np.ndarray([1, 2, 4, 8, 16, 32]) - test_cfg = CpuDetectorConfig(model=ModelConfig()) + test_cfg = DetectorConfig(type="cpu", model=ModelConfig()) test_cfg.model.input_tensor = InputTensorEnum.nchw test_obj_detect = frigate.object_detection.LocalObjectDetector( @@ -121,7 +113,7 @@ class TestLocalObjectDetector(unittest.TestCase): "label-5", ] - test_cfg = CpuDetectorConfig(model=ModelConfig()) + test_cfg = DetectorConfig(type="cpu", model=ModelConfig()) test_cfg.model = ModelConfig() test_obj_detect = frigate.object_detection.LocalObjectDetector( detector_config=test_cfg,