Refactor TFL so code is shared among all detectors (#24822)

frigate/detectors/detector_utils.py

@@ -1,8 +1,13 @@
 import logging
+import math
 import os
 
+import cv2
 import numpy as np
 
+from frigate.detectors.detector_config import ModelConfig, ModelTypeEnum
+from frigate.util.model import xyxy_to_xywh_for_nms
+
 try:
     from tflite_runtime.interpreter import Interpreter, load_delegate
 except ModuleNotFoundError:
@@ -12,39 +17,324 @@ except ModuleNotFoundError:
 logger = logging.getLogger(__name__)
 
 
-def tflite_init(self, interpreter):
-    self.interpreter = interpreter
+class TFLiteDetectionModel:
+    """Object detection pre and post processing for a TFLite interpreter.
 
-    self.interpreter.allocate_tensors()
+    Every TFLite detector shares this, differing only in how it builds the
+    interpreter (on the CPU, or with the EdgeTPU or Teflon delegate).
+    """
 
-    self.tensor_input_details = self.interpreter.get_input_details()
-    self.tensor_output_details = self.interpreter.get_output_details()
+    supported_models = [
+        ModelTypeEnum.ssd,
+        ModelTypeEnum.yologeneric,
+    ]
 
+    def __init__(
+        self, interpreter: Interpreter, model_config: ModelConfig, detector_key: str
+    ):
+        """Prepare the interpreter and work out the model's output layout.
 
-def tflite_detect_raw(self, tensor_input):
-    self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
-    self.interpreter.invoke()
+        Args:
+            interpreter: The TFLite interpreter, with any delegate loaded
+            model_config: The model being run
+            detector_key: The detector's type, for error messages
+        """
+        self.model_type = model_config.model_type
 
-    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])
+        if self.model_type not in self.supported_models:
+            raise ValueError(
+                f"{self.model_type.value} is currently not supported for {detector_key}. See the docs for more info on supported models."
+            )
 
-    detections = np.zeros((20, 6), np.float32)
+        self.interpreter = interpreter
 
-    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],
-        ]
+        try:
+            self.interpreter.allocate_tensors()
+        except RuntimeError as e:
+            # an EdgeTPU compiled model wraps its graph in a custom op that only
+            # the EdgeTPU delegate can run
+            if "edgetpu-custom-op" in str(e):
+                logger.error(
+                    "%s is compiled for the EdgeTPU and can only be run by an edgetpu detector",
+                    model_config.path,
+                )
+            raise
 
-    return detections
+        self.tensor_input_details = self.interpreter.get_input_details()
+        self.tensor_output_details = self.interpreter.get_output_details()
+        self.model_width = model_config.width
+        self.model_height = model_config.height
+
+        self.min_score = 0.4
+        self.max_detections = 20
+
+        self.model_requires_int8 = self.tensor_input_details[0]["dtype"] == np.int8
+
+        if self.model_type == ModelTypeEnum.yologeneric:
+            logger.debug("Using YOLO preprocessing/postprocessing")
+
+            if len(self.tensor_output_details) not in [2, 3]:
+                raise ValueError(
+                    f"Invalid count of output tensors in YOLO model. Found {len(self.tensor_output_details)}, expecting 2 or 3."
+                )
+
+            self.reg_max = 16  # = 64 dfl_channels // 4 # YOLO standard
+            self.min_logit_value = np.log(
+                self.min_score / (1 - self.min_score)
+            )  # for filtering
+            self._generate_anchors_and_strides()  # decode bounding box DFL
+            self.project = np.arange(
+                self.reg_max, dtype=np.float32
+            )  # for decoding bounding box DFL information
+
+            # Determine YOLO tensor indices and quantization scales for
+            # boxes and class_scores the tensor ordering and names are
+            # not reliable, so use tensor shape to detect which tensor
+            # holds boxes or class scores.
+            # The tensors have shapes (B, N, C)
+            # where N is the number of candidates (=2100 for 320x320)
+            # this may guess wrong if the number of classes is exactly 64
+            output_boxes_index = None
+            output_classes_index = None
+            for i, x in enumerate(self.tensor_output_details):
+                # the nominal index seems to start at 1 instead of 0
+                if len(x["shape"]) == 3 and x["shape"][2] == 64:
+                    output_boxes_index = i
+                elif len(x["shape"]) == 3 and x["shape"][2] > 1:
+                    # require the number of classes to be more than 1
+                    # to differentiate from (not used) max score tensor
+                    output_classes_index = i
+            if output_boxes_index is None or output_classes_index is None:
+                logger.warning("Unrecognized model output, unexpected tensor shapes.")
+                output_classes_index = (
+                    0
+                    if (output_boxes_index is None or output_classes_index == 1)
+                    else 1
+                )  # 0 is default guess
+                output_boxes_index = 1 if (output_boxes_index == 0) else 0
+
+            scores_details = self.tensor_output_details[output_classes_index]
+            boxes_details = self.tensor_output_details[output_boxes_index]
+
+            # post processing works on the raw quantized values, which a float
+            # model doesn't have
+            if np.issubdtype(scores_details["dtype"], np.floating) or np.issubdtype(
+                boxes_details["dtype"], np.floating
+            ):
+                raise ValueError(
+                    f"{model_config.path} has float outputs. YOLO models for TFLite must be int8 quantized."
+                )
+
+            self.scores_tensor_index = scores_details["index"]
+            self.scores_scale, self.scores_zero_point = scores_details["quantization"]
+            # calculate the quantized version of the min_score
+            self.min_score_quantized = int(
+                (self.min_logit_value / self.scores_scale) + self.scores_zero_point
+            )
+            self.logit_shift_to_positive_values = (
+                max(0, math.ceil((128 + self.scores_zero_point) * self.scores_scale))
+                + 1
+            )  # round up
+
+            self.boxes_tensor_index = boxes_details["index"]
+            self.boxes_scale, self.boxes_zero_point = boxes_details["quantization"]
+
+        elif self.model_type == ModelTypeEnum.ssd:
+            logger.debug("Using SSD preprocessing/postprocessing")
+
+            # SSD models end in TFLite_Detection_PostProcess, which outputs
+            # boxes, class ids, scores and count
+            if len(self.tensor_output_details) < 4:
+                raise ValueError(
+                    f"Invalid count of output tensors in SSD model. Found {len(self.tensor_output_details)}, expecting 4. YOLO models need a model_type of yolo-generic."
+                )
+
+    def _generate_anchors_and_strides(self):
+        # for decoding the bounding box DFL information into xy coordinates
+        all_anchors = []
+        all_strides = []
+        strides = (8, 16, 32)  # YOLO's small, medium, large detection heads
+
+        for stride in strides:
+            feat_h, feat_w = self.model_height // stride, self.model_width // stride
+
+            grid_y, grid_x = np.meshgrid(
+                np.arange(feat_h, dtype=np.float32),
+                np.arange(feat_w, dtype=np.float32),
+                indexing="ij",
+            )
+
+            grid_coords = np.stack((grid_x.flatten(), grid_y.flatten()), axis=1)
+            anchor_points = grid_coords + 0.5
+
+            all_anchors.append(anchor_points)
+            all_strides.append(np.full((feat_h * feat_w, 1), stride, dtype=np.float32))
+
+        self.anchors = np.concatenate(all_anchors, axis=0)
+        self.anchor_strides = np.concatenate(all_strides, axis=0)
+
+    def pre_process(self, tensor_input):
+        if self.model_requires_int8:
+            tensor_input = np.bitwise_xor(tensor_input, 128).view(
+                np.int8
+            )  # shift by -128
+        return tensor_input
+
+    def detect_raw(self, tensor_input):
+        """Run the model on a frame.
+
+        Args:
+            tensor_input: The frame, already resized to the model's input
+
+        Returns:
+            Up to 20 detections as [class_id, score, ymin, xmin, ymax, xmax],
+            padded with zeros
+        """
+        tensor_input = self.pre_process(tensor_input)
+
+        self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
+        self.interpreter.invoke()
+
+        if self.model_type == ModelTypeEnum.yologeneric:
+            return self._post_process_yolo()
+
+        return self._post_process_ssd()
+
+    def _post_process_yolo(self):
+        # Multi-tensor YOLO model with (non-standard B(H*W)C output format).
+        # (the comments indicate the shape of tensors,
+        # using "2100" as the anchor count (for image size of 320x320),
+        # "NC" as number of classes,
+        # "N" as the count that survive after min-score filtering)
+        # TENSOR A) class scores (1, 2100, NC) with logit values
+        # TENSOR B) box coordinates (1, 2100, 64) encoded as dfl scores
+        # Recommend that the model clamp the logit values in tensor (A)
+        # to the range [-4,+4] to preserve precision from [2%,98%]
+        # and because NMS requires the min_score parameter to be >= 0
+
+        # don't dequantize scores data yet, wait until the low-confidence
+        # candidates are filtered out from the overall result set.
+        # This reduces the work and makes post-processing faster.
+        # this method works with raw quantized numbers when possible,
+        # which relies on the value of the scale factor to be >0.
+        # This speeds up max and argmax operations.
+        # Get max confidence for each detection and create the mask
+        detections = np.zeros(
+            (self.max_detections, 6), np.float32
+        )  # initialize zero results
+        scores_output_quantized = self.interpreter.get_tensor(self.scores_tensor_index)[
+            0
+        ]  # (2100, NC)
+        max_scores_quantized = np.max(scores_output_quantized, axis=1)  # (2100,)
+        mask = max_scores_quantized >= self.min_score_quantized  # (2100,)
+
+        if not np.any(mask):
+            return detections  # empty results
+
+        max_scores_filtered_shiftedpositive = (
+            (max_scores_quantized[mask] - self.scores_zero_point) * self.scores_scale
+        ) + self.logit_shift_to_positive_values  # (N,1) shifted logit values
+        scores_output_quantized_filtered = scores_output_quantized[mask]
+
+        # dequantize boxes. NMS needs them to be in float format
+        # remove candidates with probabilities < threshold
+        boxes_output_quantized_filtered = (
+            self.interpreter.get_tensor(self.boxes_tensor_index)[0]
+        )[mask]  # (N, 64)
+        boxes_output_filtered = (
+            boxes_output_quantized_filtered.astype(np.float32) - self.boxes_zero_point
+        ) * self.boxes_scale
+
+        # 2. Decode DFL to distances (ltrb)
+        dfl_distributions = boxes_output_filtered.reshape(
+            -1, 4, self.reg_max
+        )  # (N, 4, 16)
+
+        # Softmax over the 16 bins
+        dfl_max = np.max(dfl_distributions, axis=2, keepdims=True)
+        dfl_exp = np.exp(dfl_distributions - dfl_max)
+        dfl_probs = dfl_exp / np.sum(dfl_exp, axis=2, keepdims=True)  # (N, 4, 16)
+
+        # Weighted sum: (N, 4, 16) * (16,) -> (N, 4)
+        distances = np.einsum("pcr,r->pc", dfl_probs, self.project)
+
+        # Calculate box corners in pixel coordinates
+        anchors_filtered = self.anchors[mask]
+        anchor_strides_filtered = self.anchor_strides[mask]
+        x1y1 = (
+            anchors_filtered - distances[:, [0, 1]]
+        ) * anchor_strides_filtered  # (N, 2)
+        x2y2 = (
+            anchors_filtered + distances[:, [2, 3]]
+        ) * anchor_strides_filtered  # (N, 2)
+        boxes_filtered_decoded = np.concatenate((x1y1, x2y2), axis=-1)  # (N, 4)
+
+        # 9. Apply NMS. Use logit scores here to defer sigmoid()
+        # until after filtering out redundant boxes
+        # Shift the logit scores to be non-negative (required by cv2)
+        indices = cv2.dnn.NMSBoxes(
+            bboxes=xyxy_to_xywh_for_nms(boxes_filtered_decoded),
+            scores=max_scores_filtered_shiftedpositive,
+            score_threshold=(
+                self.min_logit_value + self.logit_shift_to_positive_values
+            ),
+            nms_threshold=0.4,  # should this be a model config setting?
+        )
+        num_detections = len(indices)
+        if num_detections == 0:
+            return detections  # empty results
+
+        nms_indices = np.array(indices, dtype=np.int32).ravel()  # or .flatten()
+        if num_detections > self.max_detections:
+            nms_indices = nms_indices[: self.max_detections]
+            num_detections = self.max_detections
+        kept_logits_quantized = scores_output_quantized_filtered[nms_indices]
+        class_ids_post_nms = np.argmax(kept_logits_quantized, axis=1)
+
+        # Extract the final boxes and scores using fancy indexing
+        final_boxes = boxes_filtered_decoded[nms_indices]
+        final_scores_logits = (
+            max_scores_filtered_shiftedpositive[nms_indices]
+            - self.logit_shift_to_positive_values
+        )  # Unshifted logits
+
+        # Detections array format: [class_id, score, ymin, xmin, ymax, xmax]
+        detections[:num_detections, 0] = class_ids_post_nms
+        detections[:num_detections, 1] = 1.0 / (
+            1.0 + np.exp(-final_scores_logits)
+        )  # sigmoid
+        detections[:num_detections, 2] = final_boxes[:, 1] / self.model_height
+        detections[:num_detections, 3] = final_boxes[:, 0] / self.model_width
+        detections[:num_detections, 4] = final_boxes[:, 3] / self.model_height
+        detections[:num_detections, 5] = final_boxes[:, 2] / self.model_width
+        return detections
+
+    def _post_process_ssd(self):
+        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((self.max_detections, 6), np.float32)
+
+        for i in range(count):
+            if scores[i] < self.min_score:
+                break
+            if i == self.max_detections:
+                logger.debug("Too many detections (%s)", count)
+                break
+            detections[i] = [
+                class_ids[i],
+                float(scores[i]),
+                boxes[i][0],
+                boxes[i][1],
+                boxes[i][2],
+                boxes[i][3],
+            ]
+
+        return detections
 
 
 def tflite_load_delegate_interpreter(

frigate/detectors/plugins/cpu_tfl.py

@@ -7,7 +7,7 @@ from frigate.detectors.detection_api import DetectionApi
 from frigate.detectors.detector_config import BaseDetectorConfig
 from frigate.log import suppress_stderr_during
 
-from ..detector_utils import tflite_detect_raw, tflite_init
+from ..detector_utils import TFLiteDetectionModel
 
 try:
     from tflite_runtime.interpreter import Interpreter
@@ -40,6 +40,7 @@ class CpuDetectorConfig(BaseDetectorConfig):
 
 class CpuTfl(DetectionApi):
     type_key = DETECTOR_KEY
+    supported_models = TFLiteDetectionModel.supported_models
 
     def __init__(self, detector_config: CpuDetectorConfig):
         # Suppress TFLite delegate creation messages that bypass Python logging
@@ -49,7 +50,9 @@ class CpuTfl(DetectionApi):
                 num_threads=detector_config.num_threads or 3,
             )
 
-        tflite_init(self, interpreter)
+        self.tflite_model = TFLiteDetectionModel(
+            interpreter, detector_config.model, DETECTOR_KEY
+        )
 
     def detect_raw(self, tensor_input):
-        return tflite_detect_raw(self, tensor_input)
+        return self.tflite_model.detect_raw(tensor_input)

frigate/detectors/plugins/edgetpu_tfl.py

@@ -1,15 +1,13 @@
 import logging
-import math
 import os
 from typing import ClassVar, Literal
 
-import cv2
-import numpy as np
 from pydantic import ConfigDict, Field
 
 from frigate.detectors.detection_api import DetectionApi
-from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
-from frigate.util.model import xyxy_to_xywh_for_nms
+from frigate.detectors.detector_config import BaseDetectorConfig
+
+from ..detector_utils import TFLiteDetectionModel
 
 try:
     from tflite_runtime.interpreter import Interpreter, load_delegate
@@ -41,10 +39,7 @@ class EdgeTpuDetectorConfig(BaseDetectorConfig):
 
 class EdgeTpuTfl(DetectionApi):
     type_key = DETECTOR_KEY
-    supported_models = [
-        ModelTypeEnum.ssd,
-        ModelTypeEnum.yologeneric,
-    ]
+    supported_models = TFLiteDetectionModel.supported_models
 
     def __init__(self, detector_config: EdgeTpuDetectorConfig):
         device_config = {}
@@ -60,7 +55,7 @@ class EdgeTpuTfl(DetectionApi):
             logger.info(f"Attempting to load TPU as {device_type}")
             edge_tpu_delegate = load_delegate("libedgetpu.so.1.0", device_config)
             logger.info("TPU found")
-            self.interpreter = Interpreter(
+            interpreter = Interpreter(
                 model_path=detector_config.model.path,
                 experimental_delegates=[edge_tpu_delegate],
             )
@@ -78,298 +73,9 @@ class EdgeTpuTfl(DetectionApi):
 
             raise
 
-        self.interpreter.allocate_tensors()
-
-        self.tensor_input_details = self.interpreter.get_input_details()
-        self.tensor_output_details = self.interpreter.get_output_details()
-        self.model_width = detector_config.model.width
-        self.model_height = detector_config.model.height
-
-        self.min_score = 0.4
-        self.max_detections = 20
-
-        self.model_type = detector_config.model.model_type
-        self.model_requires_int8 = self.tensor_input_details[0]["dtype"] == np.int8
-
-        if self.model_type == ModelTypeEnum.yologeneric:
-            logger.debug("Using YOLO preprocessing/postprocessing")
-
-            if len(self.tensor_output_details) not in [2, 3]:
-                logger.error(
-                    f"Invalid count of output tensors in YOLO model. Found {len(self.tensor_output_details)}, expecting 2 or 3."
-                )
-                raise
-
-            self.reg_max = 16  # = 64 dfl_channels // 4 # YOLO standard
-            self.min_logit_value = np.log(
-                self.min_score / (1 - self.min_score)
-            )  # for filtering
-            self._generate_anchors_and_strides()  # decode bounding box DFL
-            self.project = np.arange(
-                self.reg_max, dtype=np.float32
-            )  # for decoding bounding box DFL information
-
-            # Determine YOLO tensor indices and quantization scales for
-            # boxes and class_scores the tensor ordering and names are
-            # not reliable, so use tensor shape to detect which tensor
-            # holds boxes or class scores.
-            # The tensors have shapes (B, N, C)
-            # where N is the number of candidates (=2100 for 320x320)
-            # this may guess wrong if the number of classes is exactly 64
-            output_boxes_index = None
-            output_classes_index = None
-            for i, x in enumerate(self.tensor_output_details):
-                # the nominal index seems to start at 1 instead of 0
-                if len(x["shape"]) == 3 and x["shape"][2] == 64:
-                    output_boxes_index = i
-                elif len(x["shape"]) == 3 and x["shape"][2] > 1:
-                    # require the number of classes to be more than 1
-                    # to differentiate from (not used) max score tensor
-                    output_classes_index = i
-            if output_boxes_index is None or output_classes_index is None:
-                logger.warning("Unrecognized model output, unexpected tensor shapes.")
-                output_classes_index = (
-                    0
-                    if (output_boxes_index is None or output_classes_index == 1)
-                    else 1
-                )  # 0 is default guess
-                output_boxes_index = 1 if (output_boxes_index == 0) else 0
-
-            scores_details = self.tensor_output_details[output_classes_index]
-            self.scores_tensor_index = scores_details["index"]
-            self.scores_scale, self.scores_zero_point = scores_details["quantization"]
-            # calculate the quantized version of the min_score
-            self.min_score_quantized = int(
-                (self.min_logit_value / self.scores_scale) + self.scores_zero_point
-            )
-            self.logit_shift_to_positive_values = (
-                max(0, math.ceil((128 + self.scores_zero_point) * self.scores_scale))
-                + 1
-            )  # round up
-
-            boxes_details = self.tensor_output_details[output_boxes_index]
-            self.boxes_tensor_index = boxes_details["index"]
-            self.boxes_scale, self.boxes_zero_point = boxes_details["quantization"]
-
-        elif self.model_type == ModelTypeEnum.ssd:
-            logger.debug("Using SSD preprocessing/postprocessing")
-
-            # SSD model indices (4 outputs: boxes, class_ids, scores, count)
-            for x in self.tensor_output_details:
-                if len(x["shape"]) == 3:
-                    self.output_boxes_index = x["index"]
-                elif len(x["shape"]) == 1:
-                    self.output_count_index = x["index"]
-
-            self.output_class_ids_index = None
-            self.output_class_scores_index = None
-
-        else:
-            raise Exception(
-                f"{self.model_type} is currently not supported for edgetpu. See the docs for more info on supported models."
-            )
-
-    def _generate_anchors_and_strides(self):
-        # for decoding the bounding box DFL information into xy coordinates
-        all_anchors = []
-        all_strides = []
-        strides = (8, 16, 32)  # YOLO's small, medium, large detection heads
-
-        for stride in strides:
-            feat_h, feat_w = self.model_height // stride, self.model_width // stride
-
-            grid_y, grid_x = np.meshgrid(
-                np.arange(feat_h, dtype=np.float32),
-                np.arange(feat_w, dtype=np.float32),
-                indexing="ij",
-            )
-
-            grid_coords = np.stack((grid_x.flatten(), grid_y.flatten()), axis=1)
-            anchor_points = grid_coords + 0.5
-
-            all_anchors.append(anchor_points)
-            all_strides.append(np.full((feat_h * feat_w, 1), stride, dtype=np.float32))
-
-        self.anchors = np.concatenate(all_anchors, axis=0)
-        self.anchor_strides = np.concatenate(all_strides, axis=0)
-
-    def determine_indexes_for_non_yolo_models(self):
-        """Legacy method for SSD models."""
-        if (
-            self.output_class_ids_index is None
-            or self.output_class_scores_index is None
-        ):
-            for i in range(4):
-                index = self.tensor_output_details[i]["index"]
-                if (
-                    index != self.output_boxes_index
-                    and index != self.output_count_index
-                ):
-                    if (
-                        np.mod(np.float32(self.interpreter.tensor(index)()[0][0]), 1)
-                        == 0.0
-                    ):
-                        self.output_class_ids_index = index
-                    else:
-                        self.output_scores_index = index
-
-    def pre_process(self, tensor_input):
-        if self.model_requires_int8:
-            tensor_input = np.bitwise_xor(tensor_input, 128).view(
-                np.int8
-            )  # shift by -128
-        return tensor_input
+        self.tflite_model = TFLiteDetectionModel(
+            interpreter, detector_config.model, DETECTOR_KEY
+        )
 
     def detect_raw(self, tensor_input):
-        tensor_input = self.pre_process(tensor_input)
-
-        self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input)
-        self.interpreter.invoke()
-
-        if self.model_type == ModelTypeEnum.yologeneric:
-            # Multi-tensor YOLO model with (non-standard B(H*W)C output format).
-            # (the comments indicate the shape of tensors,
-            # using "2100" as the anchor count (for image size of 320x320),
-            # "NC" as number of classes,
-            # "N" as the count that survive after min-score filtering)
-            # TENSOR A) class scores (1, 2100, NC) with logit values
-            # TENSOR B) box coordinates (1, 2100, 64) encoded as dfl scores
-            # Recommend that the model clamp the logit values in tensor (A)
-            # to the range [-4,+4] to preserve precision from [2%,98%]
-            # and because NMS requires the min_score parameter to be >= 0
-
-            # don't dequantize scores data yet, wait until the low-confidence
-            # candidates are filtered out from the overall result set.
-            # This reduces the work and makes post-processing faster.
-            # this method works with raw quantized numbers when possible,
-            # which relies on the value of the scale factor to be >0.
-            # This speeds up max and argmax operations.
-            # Get max confidence for each detection and create the mask
-            detections = np.zeros(
-                (self.max_detections, 6), np.float32
-            )  # initialize zero results
-            scores_output_quantized = self.interpreter.get_tensor(
-                self.scores_tensor_index
-            )[0]  # (2100, NC)
-            max_scores_quantized = np.max(scores_output_quantized, axis=1)  # (2100,)
-            mask = max_scores_quantized >= self.min_score_quantized  # (2100,)
-
-            if not np.any(mask):
-                return detections  # empty results
-
-            max_scores_filtered_shiftedpositive = (
-                (max_scores_quantized[mask] - self.scores_zero_point)
-                * self.scores_scale
-            ) + self.logit_shift_to_positive_values  # (N,1) shifted logit values
-            scores_output_quantized_filtered = scores_output_quantized[mask]
-
-            # dequantize boxes. NMS needs them to be in float format
-            # remove candidates with probabilities < threshold
-            boxes_output_quantized_filtered = (
-                self.interpreter.get_tensor(self.boxes_tensor_index)[0]
-            )[mask]  # (N, 64)
-            boxes_output_filtered = (
-                boxes_output_quantized_filtered.astype(np.float32)
-                - self.boxes_zero_point
-            ) * self.boxes_scale
-
-            # 2. Decode DFL to distances (ltrb)
-            dfl_distributions = boxes_output_filtered.reshape(
-                -1, 4, self.reg_max
-            )  # (N, 4, 16)
-
-            # Softmax over the 16 bins
-            dfl_max = np.max(dfl_distributions, axis=2, keepdims=True)
-            dfl_exp = np.exp(dfl_distributions - dfl_max)
-            dfl_probs = dfl_exp / np.sum(dfl_exp, axis=2, keepdims=True)  # (N, 4, 16)
-
-            # Weighted sum: (N, 4, 16) * (16,) -> (N, 4)
-            distances = np.einsum("pcr,r->pc", dfl_probs, self.project)
-
-            # Calculate box corners in pixel coordinates
-            anchors_filtered = self.anchors[mask]
-            anchor_strides_filtered = self.anchor_strides[mask]
-            x1y1 = (
-                anchors_filtered - distances[:, [0, 1]]
-            ) * anchor_strides_filtered  # (N, 2)
-            x2y2 = (
-                anchors_filtered + distances[:, [2, 3]]
-            ) * anchor_strides_filtered  # (N, 2)
-            boxes_filtered_decoded = np.concatenate((x1y1, x2y2), axis=-1)  # (N, 4)
-
-            # 9. Apply NMS. Use logit scores here to defer sigmoid()
-            # until after filtering out redundant boxes
-            # Shift the logit scores to be non-negative (required by cv2)
-            indices = cv2.dnn.NMSBoxes(
-                bboxes=xyxy_to_xywh_for_nms(boxes_filtered_decoded),
-                scores=max_scores_filtered_shiftedpositive,
-                score_threshold=(
-                    self.min_logit_value + self.logit_shift_to_positive_values
-                ),
-                nms_threshold=0.4,  # should this be a model config setting?
-            )
-            num_detections = len(indices)
-            if num_detections == 0:
-                return detections  # empty results
-
-            nms_indices = np.array(indices, dtype=np.int32).ravel()  # or .flatten()
-            if num_detections > self.max_detections:
-                nms_indices = nms_indices[: self.max_detections]
-                num_detections = self.max_detections
-            kept_logits_quantized = scores_output_quantized_filtered[nms_indices]
-            class_ids_post_nms = np.argmax(kept_logits_quantized, axis=1)
-
-            # Extract the final boxes and scores using fancy indexing
-            final_boxes = boxes_filtered_decoded[nms_indices]
-            final_scores_logits = (
-                max_scores_filtered_shiftedpositive[nms_indices]
-                - self.logit_shift_to_positive_values
-            )  # Unshifted logits
-
-            # Detections array format: [class_id, score, ymin, xmin, ymax, xmax]
-            detections[:num_detections, 0] = class_ids_post_nms
-            detections[:num_detections, 1] = 1.0 / (
-                1.0 + np.exp(-final_scores_logits)
-            )  # sigmoid
-            detections[:num_detections, 2] = final_boxes[:, 1] / self.model_height
-            detections[:num_detections, 3] = final_boxes[:, 0] / self.model_width
-            detections[:num_detections, 4] = final_boxes[:, 3] / self.model_height
-            detections[:num_detections, 5] = final_boxes[:, 2] / self.model_width
-            return detections
-
-        elif self.model_type == ModelTypeEnum.ssd:
-            self.determine_indexes_for_non_yolo_models()
-            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((self.max_detections, 6), np.float32)
-
-            for i in range(count):
-                if scores[i] < self.min_score:
-                    break
-                if i == self.max_detections:
-                    logger.debug(f"Too many detections ({count})!")
-                    break
-                detections[i] = [
-                    class_ids[i],
-                    float(scores[i]),
-                    boxes[i][0],
-                    boxes[i][1],
-                    boxes[i][2],
-                    boxes[i][3],
-                ]
-
-            return detections
-
-        else:
-            raise Exception(
-                f"{self.model_type} is currently not supported for edgetpu. See the docs for more info on supported models."
-            )
+        return self.tflite_model.detect_raw(tensor_input)

frigate/detectors/plugins/teflon_tfl.py

@@ -7,8 +7,7 @@ from frigate.detectors.detection_api import DetectionApi
 from frigate.detectors.detector_config import BaseDetectorConfig
 
 from ..detector_utils import (
-    tflite_detect_raw,
-    tflite_init,
+    TFLiteDetectionModel,
     tflite_load_delegate_interpreter,
 )
 
@@ -30,6 +29,7 @@ class TeflonDetectorConfig(BaseDetectorConfig):
 
 class TeflonTfl(DetectionApi):
     type_key = DETECTOR_KEY
+    supported_models = TFLiteDetectionModel.supported_models
 
     def __init__(self, detector_config: TeflonDetectorConfig):
         # Location in Debian's mesa-teflon-delegate
@@ -39,7 +39,9 @@ class TeflonTfl(DetectionApi):
         interpreter = tflite_load_delegate_interpreter(
             delegate_library, detector_config, device_config
         )
-        tflite_init(self, interpreter)
+        self.tflite_model = TFLiteDetectionModel(
+            interpreter, detector_config.model, DETECTOR_KEY
+        )
 
     def detect_raw(self, tensor_input):
-        return tflite_detect_raw(self, tensor_input)
+        return self.tflite_model.detect_raw(tensor_input)

frigate/test/test_tflite_detection_model.py

@@ -0,0 +1,229 @@
+"""Tests for the TFLite detection model shared by the TFLite detectors."""
+
+import unittest
+
+import numpy as np
+
+from frigate.detectors.detector_config import ModelConfig, ModelTypeEnum
+from frigate.detectors.detector_utils import TFLiteDetectionModel
+
+
+class FakeInterpreter:
+    """Stands in for a TFLite interpreter with fixed output tensors."""
+
+    def __init__(self, outputs, input_dtype=np.uint8, allocate_error=None):
+        self.outputs = {}
+        self.output_details = []
+        self.allocate_error = allocate_error
+        self.input_details = [
+            {"index": 0, "dtype": input_dtype, "shape": np.array([1, 320, 320, 3])}
+        ]
+        self.input = None
+
+        for i, (tensor, quantization) in enumerate(outputs, start=1):
+            self.outputs[i] = tensor
+            self.output_details.append(
+                {
+                    "index": i,
+                    "dtype": tensor.dtype,
+                    "shape": np.array(tensor.shape),
+                    "quantization": quantization,
+                }
+            )
+
+    def allocate_tensors(self):
+        if self.allocate_error:
+            raise self.allocate_error
+
+    def get_input_details(self):
+        return self.input_details
+
+    def get_output_details(self):
+        return self.output_details
+
+    def set_tensor(self, index, value):
+        self.input = value
+
+    def invoke(self):
+        pass
+
+    def tensor(self, index):
+        return lambda: self.outputs[index]
+
+    def get_tensor(self, index):
+        return self.outputs[index]
+
+
+def model_with_type(model_type) -> ModelConfig:
+    return ModelConfig(
+        path="/config/model.tflite",
+        model_type=model_type,
+        labelmap_path=None,
+        labelmap={0: "person"},
+        width=320,
+        height=320,
+    )
+
+
+def ssd_outputs(scores, count=None):
+    num = len(scores)
+    boxes = np.tile(np.array([0.1, 0.2, 0.3, 0.4], np.float32), (1, num, 1))
+    class_ids = np.arange(num, dtype=np.float32).reshape(1, num)
+    scores = np.array(scores, np.float32).reshape(1, num)
+    count = np.array([num if count is None else count], np.float32)
+    return [
+        (boxes, (0.0, 0)),
+        (class_ids, (0.0, 0)),
+        (scores, (0.0, 0)),
+        (count, (0.0, 0)),
+    ]
+
+
+# 320x320 YOLO heads at strides 8, 16 and 32 have 40*40 + 20*20 + 10*10 anchors
+YOLO_ANCHORS = 2100
+YOLO_CLASSES = 3
+SCORES_SCALE = 1 / 32
+BOXES_SCALE = 0.1
+
+
+def yolo_outputs(candidates, dtype=np.int8):
+    """Build quantized YOLO outputs.
+
+    Args:
+        candidates: (anchor index, class id, logit, ltrb distances in strides)
+        dtype: The output tensors' type
+
+    Returns:
+        The class score and DFL box tensors with their quantization
+    """
+    scores = np.full((1, YOLO_ANCHORS, YOLO_CLASSES), -128, dtype)
+    boxes = np.full((1, YOLO_ANCHORS, 4, 16), -100, dtype)
+
+    for anchor, class_id, logit, distances in candidates:
+        scores[0, anchor, class_id] = round(logit / SCORES_SCALE)
+
+        # a single hot bin per side decodes to that bin's distance
+        for side, distance in enumerate(distances):
+            boxes[0, anchor, side, distance] = 100
+
+    boxes = boxes.reshape(1, YOLO_ANCHORS, 64)
+
+    if np.issubdtype(dtype, np.floating):
+        return [(scores, (0.0, 0)), (boxes, (0.0, 0))]
+
+    return [(scores, (SCORES_SCALE, 0)), (boxes, (BOXES_SCALE, 0))]
+
+
+class TestTFLiteDetectionModel(unittest.TestCase):
+    def test_ssd_returns_detections_above_min_score(self):
+        interpreter = FakeInterpreter(ssd_outputs([0.9, 0.5, 0.3, 0.8]))
+        model = TFLiteDetectionModel(
+            interpreter, model_with_type(ModelTypeEnum.ssd), "cpu"
+        )
+
+        detections = model.detect_raw(np.zeros((1, 320, 320, 3), np.uint8))
+
+        # stops at the first score under the minimum, as SSD scores are sorted
+        np.testing.assert_allclose(
+            detections[:2],
+            [[0, 0.9, 0.1, 0.2, 0.3, 0.4], [1, 0.5, 0.1, 0.2, 0.3, 0.4]],
+            rtol=1e-6,
+        )
+        self.assertFalse(detections[2:].any())
+
+    def test_ssd_caps_detections_at_twenty(self):
+        interpreter = FakeInterpreter(ssd_outputs([0.9] * 25))
+        model = TFLiteDetectionModel(
+            interpreter, model_with_type(ModelTypeEnum.ssd), "cpu"
+        )
+
+        detections = model.detect_raw(np.zeros((1, 320, 320, 3), np.uint8))
+
+        self.assertEqual(detections.shape, (20, 6))
+        self.assertTrue((detections[:, 1] == np.float32(0.9)).all())
+
+    def test_ssd_model_without_four_outputs_is_rejected(self):
+        interpreter = FakeInterpreter(yolo_outputs([]))
+
+        with self.assertRaisesRegex(ValueError, "yolo-generic"):
+            TFLiteDetectionModel(interpreter, model_with_type(ModelTypeEnum.ssd), "cpu")
+
+    def test_int8_input_is_shifted_from_uint8(self):
+        interpreter = FakeInterpreter(ssd_outputs([0.9], count=0), np.int8)
+        model = TFLiteDetectionModel(
+            interpreter, model_with_type(ModelTypeEnum.ssd), "cpu"
+        )
+        frame = np.zeros((1, 320, 320, 3), np.uint8)
+        frame[0, 0, 0] = [0, 128, 255]
+
+        model.detect_raw(frame)
+
+        self.assertEqual(interpreter.input.dtype, np.int8)
+        self.assertEqual(interpreter.input[0, 0, 0].tolist(), [-128, 0, 127])
+
+    def test_yolo_decodes_quantized_outputs(self):
+        # anchor 420 is row 10, column 20 of the stride 8 head, centered on
+        # (20.5, 10.5) strides. Its neighbor overlaps it and should be removed
+        # by NMS as it has a lower score.
+        interpreter = FakeInterpreter(
+            yolo_outputs([(420, 2, 2.0, (2, 3, 4, 5)), (421, 1, 1.0, (2, 3, 4, 5))])
+        )
+        model = TFLiteDetectionModel(
+            interpreter, model_with_type(ModelTypeEnum.yologeneric), "cpu"
+        )
+
+        detections = model.detect_raw(np.zeros((1, 320, 320, 3), np.uint8))
+
+        # box corners are (20.5 - 2, 10.5 - 3) and (20.5 + 4, 10.5 + 5) strides
+        np.testing.assert_allclose(
+            detections[0],
+            [2, 1 / (1 + np.exp(-2.0)), 60 / 320, 148 / 320, 124 / 320, 196 / 320],
+            rtol=1e-4,
+        )
+        self.assertFalse(detections[1:].any())
+
+    def test_yolo_without_candidates_returns_no_detections(self):
+        interpreter = FakeInterpreter(yolo_outputs([]))
+        model = TFLiteDetectionModel(
+            interpreter, model_with_type(ModelTypeEnum.yologeneric), "cpu"
+        )
+
+        detections = model.detect_raw(np.zeros((1, 320, 320, 3), np.uint8))
+
+        self.assertFalse(detections.any())
+
+    def test_yolo_model_with_float_outputs_is_rejected(self):
+        interpreter = FakeInterpreter(yolo_outputs([], np.float32))
+
+        with self.assertRaisesRegex(ValueError, "int8 quantized"):
+            TFLiteDetectionModel(
+                interpreter, model_with_type(ModelTypeEnum.yologeneric), "cpu"
+            )
+
+    def test_unsupported_model_type_is_rejected(self):
+        interpreter = FakeInterpreter(ssd_outputs([0.9]))
+
+        with self.assertRaisesRegex(ValueError, "yolox is currently not supported"):
+            TFLiteDetectionModel(
+                interpreter, model_with_type(ModelTypeEnum.yolox), "cpu"
+            )
+
+    def test_edgetpu_compiled_model_is_explained(self):
+        interpreter = FakeInterpreter(
+            ssd_outputs([0.9]),
+            allocate_error=RuntimeError(
+                "Encountered unresolved custom op: edgetpu-custom-op."
+            ),
+        )
+
+        with (
+            self.assertLogs("frigate.detectors.detector_utils", "ERROR") as logs,
+            self.assertRaises(RuntimeError),
+        ):
+            TFLiteDetectionModel(interpreter, model_with_type(ModelTypeEnum.ssd), "cpu")
+
+        self.assertIn("compiled for the EdgeTPU", logs.output[0])
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)