From 61c3d2adc0f061a3d48477594c9b50c07bdf7c80 Mon Sep 17 00:00:00 2001 From: Dennis George Date: Sat, 10 Dec 2022 06:42:57 -0600 Subject: [PATCH] enable auto-discovery of detectors --- frigate/detectors/__init__.py | 12 +-------- frigate/detectors/config.py | 2 +- frigate/detectors/detection_api.py | 6 ++--- frigate/detectors/detector_type.py | 8 ------ frigate/detectors/detector_types.py | 26 +++++++++++++++++++ frigate/detectors/plugins/__init__.py | 0 frigate/detectors/{ => plugins}/cpu_tfl.py | 7 ++--- .../detectors/{ => plugins}/edgetpu_tfl.py | 7 ++--- frigate/detectors/{ => plugins}/openvino.py | 7 ++--- frigate/object_detection.py | 2 +- 10 files changed, 44 insertions(+), 33 deletions(-) delete mode 100644 frigate/detectors/detector_type.py create mode 100644 frigate/detectors/detector_types.py create mode 100644 frigate/detectors/plugins/__init__.py rename frigate/detectors/{ => plugins}/cpu_tfl.py (91%) rename frigate/detectors/{ => plugins}/edgetpu_tfl.py (93%) rename frigate/detectors/{ => plugins}/openvino.py (92%) diff --git a/frigate/detectors/__init__.py b/frigate/detectors/__init__.py index bd93812a0..741fa1829 100644 --- a/frigate/detectors/__init__.py +++ b/frigate/detectors/__init__.py @@ -1,22 +1,12 @@ import logging -from enum import Enum from .detection_api import DetectionApi -from .detector_type import DetectorTypeEnum -from .cpu_tfl import CpuTfl -from .edgetpu_tfl import EdgeTpuTfl -from .openvino import OvDetector +from .detector_types import DetectorTypeEnum, api_types from .config import ModelConfig, DetectorConfig logger = logging.getLogger(__name__) -api_types = { - DetectorTypeEnum.cpu: CpuTfl, - DetectorTypeEnum.edgetpu: EdgeTpuTfl, - DetectorTypeEnum.openvino: OvDetector, -} - def create_detector(detector_config: DetectorConfig): if detector_config.type == DetectorTypeEnum.cpu: diff --git a/frigate/detectors/config.py b/frigate/detectors/config.py index 529f0dcc1..30c404211 100644 --- a/frigate/detectors/config.py +++ b/frigate/detectors/config.py @@ -8,7 +8,7 @@ from pydantic.fields import PrivateAttr from frigate.enums import InputTensorEnum, PixelFormatEnum from frigate.util import load_labels -from .detector_type import DetectorTypeEnum +from .detector_types import DetectorTypeEnum logger = logging.getLogger(__name__) diff --git a/frigate/detectors/detection_api.py b/frigate/detectors/detection_api.py index 244195d46..60a90b1b5 100644 --- a/frigate/detectors/detection_api.py +++ b/frigate/detectors/detection_api.py @@ -1,15 +1,15 @@ import logging - from abc import ABC, abstractmethod -from typing import Dict logger = logging.getLogger(__name__) class DetectionApi(ABC): + type_key: str + @abstractmethod - def __init__(self, det_device=None, model_config=None): + def __init__(self, detector_config): pass @abstractmethod diff --git a/frigate/detectors/detector_type.py b/frigate/detectors/detector_type.py deleted file mode 100644 index 266ae13d2..000000000 --- a/frigate/detectors/detector_type.py +++ /dev/null @@ -1,8 +0,0 @@ -from enum import Enum - - -class DetectorTypeEnum(str, Enum): - edgetpu = "edgetpu" - openvino = "openvino" - cpu = "cpu" - tensorrt = "tensorrt" diff --git a/frigate/detectors/detector_types.py b/frigate/detectors/detector_types.py new file mode 100644 index 000000000..48505c218 --- /dev/null +++ b/frigate/detectors/detector_types.py @@ -0,0 +1,26 @@ +import logging +import importlib +import pkgutil +from enum import Enum + +from . import plugins +from .detection_api import DetectionApi + + +logger = logging.getLogger(__name__) + +plugin_modules = [ + importlib.import_module(name) + for finder, name, ispkg in pkgutil.iter_modules( + plugins.__path__, plugins.__name__ + "." + ) +] + +api_types = {det.type_key: det for det in DetectionApi.__subclasses__()} + + +class StrEnum(str, Enum): + pass + + +DetectorTypeEnum = StrEnum("DetectorTypeEnum", {k: k for k in api_types}) diff --git a/frigate/detectors/plugins/__init__.py b/frigate/detectors/plugins/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/frigate/detectors/cpu_tfl.py b/frigate/detectors/plugins/cpu_tfl.py similarity index 91% rename from frigate/detectors/cpu_tfl.py rename to frigate/detectors/plugins/cpu_tfl.py index 76d3ce156..1557497e2 100644 --- a/frigate/detectors/cpu_tfl.py +++ b/frigate/detectors/plugins/cpu_tfl.py @@ -1,8 +1,7 @@ import logging import numpy as np -from .detection_api import DetectionApi -from .config import CpuDetectorConfig +from frigate.detectors.detection_api import DetectionApi import tflite_runtime.interpreter as tflite @@ -10,7 +9,9 @@ logger = logging.getLogger(__name__) class CpuTfl(DetectionApi): - def __init__(self, detector_config: CpuDetectorConfig): + type_key = "cpu" + + def __init__(self, detector_config): self.interpreter = tflite.Interpreter( model_path=detector_config.model.path or "/cpu_model.tflite", num_threads=detector_config.num_threads, diff --git a/frigate/detectors/edgetpu_tfl.py b/frigate/detectors/plugins/edgetpu_tfl.py similarity index 93% rename from frigate/detectors/edgetpu_tfl.py rename to frigate/detectors/plugins/edgetpu_tfl.py index 2994ee5b8..a08fa9e28 100644 --- a/frigate/detectors/edgetpu_tfl.py +++ b/frigate/detectors/plugins/edgetpu_tfl.py @@ -1,8 +1,7 @@ import logging import numpy as np -from .detection_api import DetectionApi -from .config import EdgeTpuDetectorConfig +from frigate.detectors.detection_api import DetectionApi import tflite_runtime.interpreter as tflite from tflite_runtime.interpreter import load_delegate @@ -11,7 +10,9 @@ logger = logging.getLogger(__name__) class EdgeTpuTfl(DetectionApi): - def __init__(self, detector_config: EdgeTpuDetectorConfig): + type_key = "edgetpu" + + def __init__(self, detector_config): device_config = {"device": "usb"} if not detector_config.device is None: device_config = {"device": detector_config.device} diff --git a/frigate/detectors/openvino.py b/frigate/detectors/plugins/openvino.py similarity index 92% rename from frigate/detectors/openvino.py rename to frigate/detectors/plugins/openvino.py index 1ed2bb405..200fffcf3 100644 --- a/frigate/detectors/openvino.py +++ b/frigate/detectors/plugins/openvino.py @@ -2,15 +2,16 @@ import logging import numpy as np import openvino.runtime as ov -from .detection_api import DetectionApi -from .config import OpenVinoDetectorConfig +from frigate.detectors.detection_api import DetectionApi logger = logging.getLogger(__name__) class OvDetector(DetectionApi): - def __init__(self, detector_config: OpenVinoDetectorConfig): + type_key = "openvino" + + def __init__(self, detector_config): self.ov_core = ov.Core() self.ov_model = self.ov_core.read_model(detector_config.model.path) diff --git a/frigate/object_detection.py b/frigate/object_detection.py index 70eaa1963..c0eb302bf 100644 --- a/frigate/object_detection.py +++ b/frigate/object_detection.py @@ -11,7 +11,7 @@ import numpy as np from setproctitle import setproctitle from frigate.enums import InputTensorEnum -from frigate.detectors import create_detector, DetectorTypeEnum +from frigate.detectors import create_detector from frigate.util import EventsPerSecond, SharedMemoryFrameManager, listen, load_labels