From bbf23ccce72948069bb05ed8489b17fa21c49d3e Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Thu, 29 Oct 2026 08:49:18 -0600 Subject: [PATCH] Refactor TFL so code is shared among all detectors (#24822) --- frigate/detectors/detector_utils.py | 340 ++++++++++++++++++-- frigate/detectors/plugins/cpu_tfl.py | 9 +- frigate/detectors/plugins/edgetpu_tfl.py | 312 +----------------- frigate/detectors/plugins/teflon_tfl.py | 10 +- frigate/test/test_tflite_detection_model.py | 229 +++++++++++++ 5 files changed, 565 insertions(+), 335 deletions(-) create mode 100644 frigate/test/test_tflite_detection_model.py diff --git a/frigate/detectors/detector_utils.py b/frigate/detectors/detector_utils.py index d8930b2ae8..a91e14068e 100644 --- a/frigate/detectors/detector_utils.py +++ b/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( diff --git a/frigate/detectors/plugins/cpu_tfl.py b/frigate/detectors/plugins/cpu_tfl.py index c33a2ba008..2f79863e9c 100644 --- a/frigate/detectors/plugins/cpu_tfl.py +++ b/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) diff --git a/frigate/detectors/plugins/edgetpu_tfl.py b/frigate/detectors/plugins/edgetpu_tfl.py index 63b51e99b8..2f213dd77f 100644 --- a/frigate/detectors/plugins/edgetpu_tfl.py +++ b/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) diff --git a/frigate/detectors/plugins/teflon_tfl.py b/frigate/detectors/plugins/teflon_tfl.py index 7dd44e6e89..8327655e9d 100644 --- a/frigate/detectors/plugins/teflon_tfl.py +++ b/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) diff --git a/frigate/test/test_tflite_detection_model.py b/frigate/test/test_tflite_detection_model.py new file mode 100644 index 0000000000..3446097448 --- /dev/null +++ b/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)