From 9394e57b311f1545967da8b12d48fd19d9fe8275 Mon Sep 17 00:00:00 2001 From: RichardPar Date: Sat, 13 May 2023 21:58:52 +0100 Subject: [PATCH] VIM4 ARMGPU tflite --- frigate/detectors/plugins/armgpu_tfl.py | 76 +++++++++++++++++++++++++ 1 file changed, 76 insertions(+) create mode 100644 frigate/detectors/plugins/armgpu_tfl.py diff --git a/frigate/detectors/plugins/armgpu_tfl.py b/frigate/detectors/plugins/armgpu_tfl.py new file mode 100644 index 000000000..c8a2591aa --- /dev/null +++ b/frigate/detectors/plugins/armgpu_tfl.py @@ -0,0 +1,76 @@ +import logging +import numpy as np + +from frigate.detectors.detection_api import DetectionApi +from frigate.detectors.detector_config import BaseDetectorConfig +from typing import Literal +from pydantic import Extra, Field + +try: + from tflite_runtime.interpreter import Interpreter, load_delegate +except ModuleNotFoundError: + from tensorflow.lite.python.interpreter import Interpreter, load_delegate + + +logger = logging.getLogger(__name__) + +DETECTOR_KEY = "armgpu" + + +class ArmGpuDetectorConfig(BaseDetectorConfig): + type: Literal[DETECTOR_KEY] + device: str = Field(default=None, title="Device Type") + + +class ArmGpuTfl(DetectionApi): + type_key = DETECTOR_KEY + + def __init__(self, detector_config: ArmGpuDetectorConfig): + + edge_tpu_delegate = None + + try: + logger.info(f"Attempting to load ARM GPU") + edge_tpu_delegate = load_delegate(library="libarmnnDelegate.so", options={"backends":"GpuAcc","logging-severity":"info"}) + logger.info("TPU found") + self.interpreter = Interpreter( + model_path=detector_config.model.path, + experimental_delegates=[edge_tpu_delegate], + ) + except ValueError: + logger.error( + "No GPU was detected. If you do not have a ARM GPU, 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