diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 3521d6f04..8b025a23c 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -26,7 +26,7 @@ RUN apt-get -qq update \ && APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=DontWarn apt-key adv --fetch-keys https://packages.cloud.google.com/apt/doc/apt-key.gpg \ && echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" > /etc/apt/sources.list.d/coral-edgetpu.list \ && echo "libedgetpu1-max libedgetpu/accepted-eula select true" | debconf-set-selections \ - && apt-get -qq update && apt-get -qq install --no-install-recommends -y libedgetpu1-max python3-tflite-runtime \ + && apt-get -qq update && apt-get -qq install --no-install-recommends -y libedgetpu1-max python3-tflite-runtime python3-pycoral \ && rm -rf /var/lib/apt/lists/* /wheels \ && (apt-get autoremove -y; apt-get autoclean -y) diff --git a/frigate/edgetpu.py b/frigate/edgetpu.py index b426e40a6..023e417b7 100644 --- a/frigate/edgetpu.py +++ b/frigate/edgetpu.py @@ -9,6 +9,7 @@ from abc import ABC, abstractmethod from typing import Dict import numpy as np +from pycoral.adapters import detect import tflite_runtime.interpreter as tflite from setproctitle import setproctitle from tflite_runtime.interpreter import load_delegate @@ -97,25 +98,20 @@ class LocalObjectDetector(ObjectDetector): def detect_raw(self, tensor_input): self.interpreter.set_tensor(self.tensor_input_details[0]["index"], tensor_input) self.interpreter.invoke() - boxes = np.squeeze( - self.interpreter.get_tensor(self.tensor_output_details[0]["index"]) - ) - label_codes = np.squeeze( - self.interpreter.get_tensor(self.tensor_output_details[1]["index"]) - ) - scores = np.squeeze( - self.interpreter.get_tensor(self.tensor_output_details[2]["index"]) - ) + + objects = detect.get_objects(self.interpreter, 0.4) detections = np.zeros((20, 6), np.float32) - for i, score in enumerate(scores): + for i, obj in enumerate(objects): + if i == 20: + break detections[i] = [ - label_codes[i], - score, - boxes[i][0], - boxes[i][1], - boxes[i][2], - boxes[i][3], + obj.id, + obj.score, + obj.bbox.ymin, + obj.bbox.xmin, + obj.bbox.ymax, + obj.bbox.xmax, ] return detections diff --git a/frigate/video.py b/frigate/video.py index 7285dc0e5..c5e6b41ce 100755 --- a/frigate/video.py +++ b/frigate/video.py @@ -404,10 +404,10 @@ def detect( for d in region_detections: box = d[2] size = region[2] - region[0] - x_min = int((box[1] * size) + region[0]) - y_min = int((box[0] * size) + region[1]) - x_max = int((box[3] * size) + region[0]) - y_max = int((box[2] * size) + region[1]) + x_min = int((box[1]) + region[0]) + y_min = int((box[0]) + region[1]) + x_max = int((box[3]) + region[0]) + y_max = int((box[2]) + region[1]) det = ( d[0], d[1],