From 80ae30e30ed1c1cde96d186485c2f02140791d68 Mon Sep 17 00:00:00 2001 From: Indrek Mandre Date: Fri, 2 Feb 2024 17:21:42 +0200 Subject: [PATCH] detectors/edgetpu_tfl: add support for yolov8 --- frigate/detectors/plugins/edgetpu_tfl.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/frigate/detectors/plugins/edgetpu_tfl.py b/frigate/detectors/plugins/edgetpu_tfl.py index ac67626a2..07dfc127d 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 +import frigate.detectors.yolo_utils as yolo_utils try: from tflite_runtime.interpreter import Interpreter, load_delegate @@ -54,11 +55,25 @@ 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 yolo_utils.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]