From 3fa174eca4ece0b63eae3722acb782460cb4f1fa Mon Sep 17 00:00:00 2001 From: Dennis George Date: Thu, 8 Dec 2022 11:25:46 -0600 Subject: [PATCH] move create_detector and DetectorTypeEnum --- frigate/detectors/__init__.py | 33 +++++++++++++++++++++++++--- frigate/detectors/cpu_tfl.py | 4 ---- frigate/detectors/detection_api.py | 15 ------------- frigate/detectors/detector_type.py | 8 ------- frigate/detectors/edgetpu_tfl.py | 4 ---- frigate/detectors/openvino.py | 4 ---- frigate/detectors/tensorrt.py | 19 ---------------- frigate/object_detection.py | 9 ++------ frigate/test/test_object_detector.py | 10 ++++----- 9 files changed, 37 insertions(+), 69 deletions(-) delete mode 100644 frigate/detectors/detector_type.py delete mode 100644 frigate/detectors/tensorrt.py diff --git a/frigate/detectors/__init__.py b/frigate/detectors/__init__.py index b8d1364e5..f5c1fb219 100644 --- a/frigate/detectors/__init__.py +++ b/frigate/detectors/__init__.py @@ -1,8 +1,35 @@ -import os +import logging +from enum import Enum -from .detector_type import DetectorTypeEnum from .detection_api import DetectionApi from .cpu_tfl import CpuTfl from .edgetpu_tfl import EdgeTpuTfl from .openvino import OvDetector -from .tensorrt import TensorRT + + +logger = logging.getLogger(__name__) + + +class DetectorTypeEnum(str, Enum): + edgetpu = "edgetpu" + openvino = "openvino" + cpu = "cpu" + tensorrt = "tensorrt" + + +def create_detector(det_type: DetectorTypeEnum, **kwargs): + _api_types = { + DetectorTypeEnum.cpu: CpuTfl, + DetectorTypeEnum.edgetpu: EdgeTpuTfl, + DetectorTypeEnum.openvino: OvDetector + } + + if det_type == DetectorTypeEnum.cpu: + logger.warning( + "CPU detectors are not recommended and should only be used for testing or for trial purposes." + ) + + api = _api_types.get(det_type) + if not api: + raise ValueError(det_type) + return api(**kwargs) diff --git a/frigate/detectors/cpu_tfl.py b/frigate/detectors/cpu_tfl.py index 0828bb38f..75896ca59 100644 --- a/frigate/detectors/cpu_tfl.py +++ b/frigate/detectors/cpu_tfl.py @@ -2,7 +2,6 @@ import logging import numpy as np from .detection_api import DetectionApi -from .detector_type import DetectorTypeEnum import tflite_runtime.interpreter as tflite @@ -46,6 +45,3 @@ class CpuTfl(DetectionApi): ] return detections - - -DetectionApi.register_api(DetectorTypeEnum.cpu, CpuTfl) diff --git a/frigate/detectors/detection_api.py b/frigate/detectors/detection_api.py index 11babed18..244195d46 100644 --- a/frigate/detectors/detection_api.py +++ b/frigate/detectors/detection_api.py @@ -1,7 +1,5 @@ import logging -from .detector_type import DetectorTypeEnum - from abc import ABC, abstractmethod from typing import Dict @@ -10,8 +8,6 @@ logger = logging.getLogger(__name__) class DetectionApi(ABC): - _api_types = {} - @abstractmethod def __init__(self, det_device=None, model_config=None): pass @@ -19,14 +15,3 @@ class DetectionApi(ABC): @abstractmethod def detect_raw(self, tensor_input): pass - - @staticmethod - def register_api(det_type: DetectorTypeEnum, det_api): - DetectionApi._api_types[det_type] = det_api - - @staticmethod - def create(det_type: DetectorTypeEnum, **kwargs): - api = DetectionApi._api_types.get(det_type) - if not api: - raise ValueError(det_type) - return api(**kwargs) 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/edgetpu_tfl.py b/frigate/detectors/edgetpu_tfl.py index 1592bccf2..c676af2fa 100644 --- a/frigate/detectors/edgetpu_tfl.py +++ b/frigate/detectors/edgetpu_tfl.py @@ -2,7 +2,6 @@ import logging import numpy as np from .detection_api import DetectionApi -from .detector_type import DetectorTypeEnum import tflite_runtime.interpreter as tflite from tflite_runtime.interpreter import load_delegate @@ -63,6 +62,3 @@ class EdgeTpuTfl(DetectionApi): ] return detections - - -DetectionApi.register_api(DetectorTypeEnum.edgetpu, EdgeTpuTfl) diff --git a/frigate/detectors/openvino.py b/frigate/detectors/openvino.py index 659f4decc..09080852b 100644 --- a/frigate/detectors/openvino.py +++ b/frigate/detectors/openvino.py @@ -3,7 +3,6 @@ import numpy as np import openvino.runtime as ov from .detection_api import DetectionApi -from .detector_type import DetectorTypeEnum logger = logging.getLogger(__name__) @@ -53,6 +52,3 @@ class OvDetector(DetectionApi): i += 1 return detections - - -DetectionApi.register_api(DetectorTypeEnum.openvino, OvDetector) diff --git a/frigate/detectors/tensorrt.py b/frigate/detectors/tensorrt.py deleted file mode 100644 index 5c6fbe54b..000000000 --- a/frigate/detectors/tensorrt.py +++ /dev/null @@ -1,19 +0,0 @@ -import logging -import numpy as np - -from .detection_api import DetectionApi -from .detector_type import DetectorTypeEnum - - -logger = logging.getLogger(__name__) - - -class TensorRT(DetectionApi): - def __init__(self, det_device=None, model_config=None, num_threads=1, **kwargs): - raise NotImplementedError("TensorRT engine is not yet functional!") - - def detect_raw(self, tensor_input): - raise NotImplementedError("TensorRT engine is not yet functional!") - - -DetectionApi.register_api(DetectorTypeEnum.tensorrt, TensorRT) diff --git a/frigate/object_detection.py b/frigate/object_detection.py index e79a94095..d67417fb6 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.config import InputTensorEnum -from frigate.detectors import DetectionApi, DetectorTypeEnum +from frigate.detectors import create_detector, DetectorTypeEnum from frigate.util import EventsPerSecond, SharedMemoryFrameManager, listen, load_labels @@ -52,12 +52,7 @@ class LocalObjectDetector(ObjectDetector): else: self.input_transform = None - if det_type == DetectorTypeEnum.cpu: - logger.warning( - "CPU detectors are not recommended and should only be used for testing or for trial purposes." - ) - - self.detect_api = DetectionApi.create(det_type, det_device=det_device, model_config=model_config, num_threads=num_threads) + self.detect_api = create_detector(det_type, det_device=det_device, model_config=model_config, num_threads=num_threads) def detect(self, tensor_input, threshold=0.4): detections = [] diff --git a/frigate/test/test_object_detector.py b/frigate/test/test_object_detector.py index ea2479284..b600360da 100644 --- a/frigate/test/test_object_detector.py +++ b/frigate/test/test_object_detector.py @@ -7,8 +7,8 @@ import frigate.object_detection class TestLocalObjectDetector(unittest.TestCase): - @patch("frigate.detectors.EdgeTpuTfl") - @patch("frigate.detectors.CpuTfl") + @patch("frigate.detectors.edgetpu_tfl.EdgeTpuTfl") + @patch("frigate.detectors.cpu_tfl.CpuTfl") def test_localdetectorprocess_given_type_cpu_should_call_cputfl_init( self, mock_cputfl, mock_edgetputfl ): @@ -22,8 +22,8 @@ class TestLocalObjectDetector(unittest.TestCase): mock_edgetputfl.assert_not_called() mock_cputfl.assert_called_once_with(model_config=test_cfg, num_threads=6) - @patch("frigate.detectors.EdgeTpuTfl") - @patch("frigate.detectors.CpuTfl") + @patch("frigate.detectors.edgetpu_tfl.EdgeTpuTfl") + @patch("frigate.detectors.cpu_tfl.CpuTfl") def test_localdetectorprocess_given_type_edgtpu_should_call_edgtpu_init( self, mock_cputfl, mock_edgetputfl ): @@ -40,7 +40,7 @@ class TestLocalObjectDetector(unittest.TestCase): mock_cputfl.assert_not_called() mock_edgetputfl.assert_called_once_with(det_device="usb", model_config=test_cfg) - @patch("frigate.detectors.CpuTfl") + @patch("frigate.detectors.cpu_tfl.CpuTfl") def test_detect_raw_given_tensor_input_should_return_api_detect_raw_result( self, mock_cputfl ):