Refactor to simplify support for additional detector types (#3656)

* Refactor EdgeTPU and CPU model handling to detector submodules.

* Fix selecting the correct detection device type from the config

* Remove detector type check when creating ObjectDetectProcess

* Fixes after rebasing to 0.11

* Add init file to detector folder

* Rename to detect_api

Co-authored-by: Nicolas Mowen <[email protected]>

* Add unit test for LocalObjectDetector class

* Add configuration for model inputs
Support transforming detection regions to RGB or BGR.
Support specifying the input tensor shape.  The tensor shape has a standard format ["BHWC"] when handed to the detector, but can be transformed in the detector to match the model shape using the model  input_tensor config.

* Add documentation for new model config parameters

* Add input tensor transpose to LocalObjectDetector

* Change the model input tensor config to use an enumeration

* Updates for model config documentation

Co-authored-by: Nicolas Mowen <[email protected]>

benchmark.py

@@ -3,10 +3,16 @@
 import multiprocessing as mp
 import numpy as np
 import datetime
-from frigate.edgetpu import LocalObjectDetector, EdgeTPUProcess, RemoteObjectDetector, load_labels
+from frigate.config import DetectorTypeEnum
+from frigate.object_detection import (
+    LocalObjectDetector,
+    ObjectDetectProcess,
+    RemoteObjectDetector,
+    load_labels,
+)
 
-my_frame = np.expand_dims(np.full((300,300,3), 1, np.uint8), axis=0)
-labels = load_labels('/labelmap.txt')
+my_frame = np.expand_dims(np.full((300, 300, 3), 1, np.uint8), axis=0)
+labels = load_labels("/labelmap.txt")
 
 ######
 # Minimal same process runner
@@ -39,20 +45,23 @@
 
 
 def start(id, num_detections, detection_queue, event):
-  object_detector = RemoteObjectDetector(str(id), '/labelmap.txt', detection_queue, event)
-  start = datetime.datetime.now().timestamp()
+    object_detector = RemoteObjectDetector(
+        str(id), "/labelmap.txt", detection_queue, event
+    )
+    start = datetime.datetime.now().timestamp()
 
-  frame_times = []
-  for x in range(0, num_detections):
-    start_frame = datetime.datetime.now().timestamp()
-    detections = object_detector.detect(my_frame)
-    frame_times.append(datetime.datetime.now().timestamp()-start_frame)
+    frame_times = []
+    for x in range(0, num_detections):
+        start_frame = datetime.datetime.now().timestamp()
+        detections = object_detector.detect(my_frame)
+        frame_times.append(datetime.datetime.now().timestamp() - start_frame)
+
+    duration = datetime.datetime.now().timestamp() - start
+    object_detector.cleanup()
+    print(f"{id} - Processed for {duration:.2f} seconds.")
+    print(f"{id} - FPS: {object_detector.fps.eps():.2f}")
+    print(f"{id} - Average frame processing time: {mean(frame_times)*1000:.2f}ms")
 
-  duration = datetime.datetime.now().timestamp()-start
-  object_detector.cleanup()
-  print(f"{id} - Processed for {duration:.2f} seconds.")
-  print(f"{id} - FPS: {object_detector.fps.eps():.2f}")
-  print(f"{id} - Average frame processing time: {mean(frame_times)*1000:.2f}ms")
 
 ######
 # Separate process runner
@@ -71,23 +80,29 @@ def start(id, num_detections, detection_queue, event):
 
 events = {}
 for x in range(0, 10):
-  events[str(x)] = mp.Event()
+    events[str(x)] = mp.Event()
 detection_queue = mp.Queue()
-edgetpu_process_1 = EdgeTPUProcess(detection_queue, events, 'usb:0')
-edgetpu_process_2 = EdgeTPUProcess(detection_queue, events, 'usb:1')
+edgetpu_process_1 = ObjectDetectProcess(
+    detection_queue, events, DetectorTypeEnum.edgetpu, "usb:0"
+)
+edgetpu_process_2 = ObjectDetectProcess(
+    detection_queue, events, DetectorTypeEnum.edgetpu, "usb:1"
+)
 
 for x in range(0, 10):
-  camera_process = mp.Process(target=start, args=(x, 300, detection_queue, events[str(x)]))
-  camera_process.daemon = True
-  camera_processes.append(camera_process)
+    camera_process = mp.Process(
+        target=start, args=(x, 300, detection_queue, events[str(x)])
+    )
+    camera_process.daemon = True
+    camera_processes.append(camera_process)
 
 start_time = datetime.datetime.now().timestamp()
 
 for p in camera_processes:
-  p.start()
+    p.start()
 
 for p in camera_processes:
-  p.join()
+    p.join()
 
-duration = datetime.datetime.now().timestamp()-start_time
-print(f"Total - Processed for {duration:.2f} seconds.")
+duration = datetime.datetime.now().timestamp() - start_time
+print(f"Total - Processed for {duration:.2f} seconds.")

docs/docs/configuration/advanced.md

@@ -23,7 +23,7 @@ Examples of available modules are:
 
 - `frigate.app`
 - `frigate.mqtt`
-- `frigate.edgetpu`
+- `frigate.object_detection`
 - `frigate.zeroconf`
 - `detector.<detector_name>`
 - `watchdog.<camera_name>`
@@ -50,6 +50,30 @@ database:
 
 If using a custom model, the width and height will need to be specified.
 
+Custom models may also require different input tensor formats. The colorspace conversion supports RGB, BGR, or YUV frames to be sent to the object detector. The input tensor shape parameter is an enumeration to match what specified by the model.
+
+| Tensor Dimension | Description    |
+| :--------------: | -------------- |
+|        N         | Batch Size     |
+|        H         | Model Height   |
+|        W         | Model Width    |
+|        C         | Color Channels |
+
+| Available Input Tensor Shapes |
+| :---------------------------: |
+|            "nhwc"             |
+|            "nchw"             |
+
+```yaml
+# Optional: model config
+model:
+  path: /path/to/model
+  width: 320
+  height: 320
+  input_tensor: "nhwc"
+  input_pixel_format: "bgr"
+```
+
 The labelmap can be customized to your needs. A common reason to do this is to combine multiple object types that are easily confused when you don't need to be as granular such as car/truck. By default, truck is renamed to car because they are often confused. You cannot add new object types, but you can change the names of existing objects in the model.
 
 ```yaml
@@ -71,6 +95,7 @@ Note that if you rename objects in the labelmap, you will also need to update yo
 Included with Frigate is a build of ffmpeg that works for the vast majority of users. However, there exists some hardware setups which have incompatibilities with the included build. In this case, a docker volume mapping can be used to overwrite the included ffmpeg build with an ffmpeg build that works for your specific hardware setup.
 
 To do this:
+
 1. Download your ffmpeg build and uncompress to a folder on the host (let's use `/home/appdata/frigate/custom-ffmpeg` for this example).
 2. Update your docker-compose or docker CLI to include `'/home/appdata/frigate/custom-ffmpeg':'/usr/lib/btbn-ffmpeg':'ro'` in the volume mappings.
 3. Restart frigate and the custom version will be used if the mapping was done correctly.

docs/docs/configuration/index.md

@@ -97,6 +97,12 @@ model:
   width: 320
   # Required: Object detection model input height (default: shown below)
   height: 320
+  # Optional: Object detection model input colorspace
+  # Valid values are rgb, bgr, or yuv. (default: shown below)
+  input_pixel_format: rgb
+  # Optional: Object detection model input tensor format
+  # Valid values are nhwc or nchw (default: shown below)
+  input_tensor: "nhwc"
   # Optional: Label name modifications. These are merged into the standard labelmap.
   labelmap:
     2: vehicle

frigate/app.py

@@ -15,7 +15,7 @@
 
 from frigate.config import DetectorTypeEnum, FrigateConfig
 from frigate.const import CACHE_DIR, CLIPS_DIR, RECORD_DIR
-from frigate.edgetpu import EdgeTPUProcess
+from frigate.object_detection import ObjectDetectProcess
 from frigate.events import EventCleanup, EventProcessor
 from frigate.http import create_app
 from frigate.log import log_process, root_configurer
@@ -40,7 +40,7 @@ class FrigateApp:
     def __init__(self) -> None:
         self.stop_event: Event = mp.Event()
         self.detection_queue: Queue = mp.Queue()
-        self.detectors: dict[str, EdgeTPUProcess] = {}
+        self.detectors: dict[str, ObjectDetectProcess] = {}
         self.detection_out_events: dict[str, Event] = {}
         self.detection_shms: list[mp.shared_memory.SharedMemory] = []
         self.log_queue: Queue = mp.Queue()
@@ -178,8 +178,6 @@ def start_mqtt_relay(self) -> None:
         self.mqtt_relay.start()
 
     def start_detectors(self) -> None:
-        model_path = self.config.model.path
-        model_shape = (self.config.model.height, self.config.model.width)
         for name in self.config.cameras.keys():
             self.detection_out_events[name] = mp.Event()
 
@@ -203,26 +201,15 @@ def start_detectors(self) -> None:
             self.detection_shms.append(shm_out)
 
         for name, detector in self.config.detectors.items():
-            if detector.type == DetectorTypeEnum.cpu:
-                self.detectors[name] = EdgeTPUProcess(
-                    name,
-                    self.detection_queue,
-                    self.detection_out_events,
-                    model_path,
-                    model_shape,
-                    "cpu",
-                    detector.num_threads,
-                )
-            if detector.type == DetectorTypeEnum.edgetpu:
-                self.detectors[name] = EdgeTPUProcess(
-                    name,
-                    self.detection_queue,
-                    self.detection_out_events,
-                    model_path,
-                    model_shape,
-                    detector.device,
-                    detector.num_threads,
-                )
+            self.detectors[name] = ObjectDetectProcess(
+                name,
+                self.detection_queue,
+                self.detection_out_events,
+                self.config.model,
+                detector.type,
+                detector.device,
+                detector.num_threads,
+            )
 
     def start_detected_frames_processor(self) -> None:
         self.detected_frames_processor = TrackedObjectProcessor(
@@ -253,7 +240,6 @@ def start_video_output_processor(self) -> None:
         logger.info(f"Output process started: {output_processor.pid}")
 
     def start_camera_processors(self) -> None:
-        model_shape = (self.config.model.height, self.config.model.width)
         for name, config in self.config.cameras.items():
             if not self.config.cameras[name].enabled:
                 logger.info(f"Camera processor not started for disabled camera {name}")
@@ -265,7 +251,7 @@ def start_camera_processors(self) -> None:
                 args=(
                     name,
                     config,
-                    model_shape,
+                    self.config.model,
                     self.config.model.merged_labelmap,
                     self.detection_queue,
                     self.detection_out_events[name],

frigate/config.py

@@ -718,6 +718,17 @@ class DatabaseConfig(FrigateBaseModel):
     )
 
 
+class PixelFormatEnum(str, Enum):
+    rgb = "rgb"
+    bgr = "bgr"
+    yuv = "yuv"
+
+
+class InputTensorEnum(str, Enum):
+    nchw = "nchw"
+    nhwc = "nhwc"
+
+
 class ModelConfig(FrigateBaseModel):
     path: Optional[str] = Field(title="Custom Object detection model path.")
     labelmap_path: Optional[str] = Field(title="Label map for custom object detector.")
@@ -726,6 +737,12 @@ class ModelConfig(FrigateBaseModel):
     labelmap: Dict[int, str] = Field(
         default_factory=dict, title="Labelmap customization."
     )
+    input_tensor: InputTensorEnum = Field(
+        default=InputTensorEnum.nhwc, title="Model Input Tensor Shape"
+    )
+    input_pixel_format: PixelFormatEnum = Field(
+        default=PixelFormatEnum.rgb, title="Model Input Pixel Color Format"
+    )
     _merged_labelmap: Optional[Dict[int, str]] = PrivateAttr()
     _colormap: Dict[int, Tuple[int, int, int]] = PrivateAttr()
 

frigate/detectors/cpu_tfl.py

@@ -0,0 +1,46 @@
+import logging
+import numpy as np
+
+from frigate.detectors.detection_api import DetectionApi
+import tflite_runtime.interpreter as tflite
+
+logger = logging.getLogger(__name__)
+
+
+class CpuTfl(DetectionApi):
+    def __init__(self, det_device=None, model_config=None, num_threads=3):
+        self.interpreter = tflite.Interpreter(
+            model_path=model_config.path or "/cpu_model.tflite", num_threads=num_threads
+        )
+
+        self.interpreter.allocate_tensors()
+
+        self.tensor_input_details = self.interpreter.get_input_details()
+        self.tensor_output_details = self.interpreter.get_output_details()
+
+    def detect_raw(self, tensor_input):
+        self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
+        self.interpreter.invoke()
+
+        boxes = self.interpreter.tensor(self.tensor_output_details[0]["index"])()[0]
+        class_ids = self.interpreter.tensor(self.tensor_output_details[1]["index"])()[0]
+        scores = self.interpreter.tensor(self.tensor_output_details[2]["index"])()[0]
+        count = int(
+            self.interpreter.tensor(self.tensor_output_details[3]["index"])()[0]
+        )
+
+        detections = np.zeros((20, 6), np.float32)
+
+        for i in range(count):
+            if scores[i] < 0.4 or i == 20:
+                break
+            detections[i] = [
+                class_ids[i],
+                float(scores[i]),
+                boxes[i][0],
+                boxes[i][1],
+                boxes[i][2],
+                boxes[i][3],
+            ]
+
+        return detections

frigate/detectors/detection_api.py

@@ -0,0 +1,17 @@
+import logging
+
+from abc import ABC, abstractmethod
+from typing import Dict
+
+
+logger = logging.getLogger(__name__)
+
+
+class DetectionApi(ABC):
+    @abstractmethod
+    def __init__(self, det_device=None, model_config=None):
+        pass
+
+    @abstractmethod
+    def detect_raw(self, tensor_input):
+        pass

frigate/detectors/edgetpu_tfl.py

@@ -0,0 +1,63 @@
+import logging
+import numpy as np
+
+from frigate.detectors.detection_api import DetectionApi
+import tflite_runtime.interpreter as tflite
+from tflite_runtime.interpreter import load_delegate
+
+logger = logging.getLogger(__name__)
+
+
+class EdgeTpuTfl(DetectionApi):
+    def __init__(self, det_device=None, model_config=None):
+        device_config = {"device": "usb"}
+        if not det_device is None:
+            device_config = {"device": det_device}
+
+        edge_tpu_delegate = None
+
+        try:
+            logger.info(f"Attempting to load TPU as {device_config['device']}")
+            edge_tpu_delegate = load_delegate("libedgetpu.so.1.0", device_config)
+            logger.info("TPU found")
+            self.interpreter = tflite.Interpreter(
+                model_path=model_config.path or "/edgetpu_model.tflite",
+                experimental_delegates=[edge_tpu_delegate],
+            )
+        except ValueError:
+            logger.error(
+                "No EdgeTPU was detected. If you do not have a Coral device yet, you must configure CPU detectors."
+            )
+            raise
+
+        self.interpreter.allocate_tensors()
+
+        self.tensor_input_details = self.interpreter.get_input_details()
+        self.tensor_output_details = self.interpreter.get_output_details()
+
+    def detect_raw(self, tensor_input):
+        self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
+        self.interpreter.invoke()
+
+        boxes = self.interpreter.tensor(self.tensor_output_details[0]["index"])()[0]
+        class_ids = self.interpreter.tensor(self.tensor_output_details[1]["index"])()[0]
+        scores = self.interpreter.tensor(self.tensor_output_details[2]["index"])()[0]
+        count = int(
+            self.interpreter.tensor(self.tensor_output_details[3]["index"])()[0]
+        )
+
+        detections = np.zeros((20, 6), np.float32)
+
+        for i in range(count):
+            if scores[i] < 0.4 or i == 20:
+                break
+            detections[i] = [
+                class_ids[i],
+                float(scores[i]),
+                boxes[i][0],
+                boxes[i][1],
+                boxes[i][2],
+                boxes[i][3],
+            ]
+
+        return detections