From e9a9467b369d8d0f29740eac172cea70a963e625 Mon Sep 17 00:00:00 2001 From: Indrek Mandre Date: Sat, 10 Feb 2024 17:56:46 +0200 Subject: [PATCH] detectors/edgetpu: add support for yolov8 models --- frigate/detectors/plugins/edgetpu_tfl.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/frigate/detectors/plugins/edgetpu_tfl.py b/frigate/detectors/plugins/edgetpu_tfl.py index ac67626a2..247e30fc8 100644 --- a/frigate/detectors/plugins/edgetpu_tfl.py +++ b/frigate/detectors/plugins/edgetpu_tfl.py @@ -6,6 +6,7 @@ from typing_extensions import Literal from frigate.detectors.detection_api import DetectionApi from frigate.detectors.detector_config import BaseDetectorConfig +from frigate.detectors.util import yolov8_postprocess try: from tflite_runtime.interpreter import Interpreter, load_delegate @@ -54,11 +55,29 @@ class EdgeTpuTfl(DetectionApi): self.tensor_input_details = self.interpreter.get_input_details() self.tensor_output_details = self.interpreter.get_output_details() + self.model_type = detector_config.model.model_type def detect_raw(self, tensor_input): + if self.model_type == "yolov8": + scale, zero_point = self.tensor_input_details[0]["quantization"] + tensor_input = ( + (tensor_input - scale * zero_point * 255) * (1.0 / (scale * 255)) + ).astype(self.tensor_input_details[0]["dtype"]) + self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input) self.interpreter.invoke() + if self.model_type == "yolov8": + scale, zero_point = self.tensor_output_details[0]["quantization"] + tensor_output = self.interpreter.get_tensor( + self.tensor_output_details[0]["index"] + ) + tensor_output = (tensor_output.astype(np.float32) - zero_point) * scale + model_input_shape = self.tensor_input_details[0]["shape"] + tensor_output[:, [0, 2]] *= model_input_shape[2] + tensor_output[:, [1, 3]] *= model_input_shape[1] + return yolov8_postprocess(model_input_shape, tensor_output) + 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]