enable auto-discovery of detectors

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:

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__)

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

frigate/detectors/detector_type.py

@@ -1,8 +0,0 @@
-from enum import Enum
-
-
-class DetectorTypeEnum(str, Enum):
-    edgetpu = "edgetpu"
-    openvino = "openvino"
-    cpu = "cpu"
-    tensorrt = "tensorrt"

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})

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,

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}

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)
 

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