Browse Source

use new pycoral libraries

Blake Blackshear 3 years ago
parent
commit
dc759a3e56
3 changed files with 17 additions and 21 deletions
  1. 1 1
      docker/Dockerfile.base
  2. 12 16
      frigate/edgetpu.py
  3. 4 4
      frigate/video.py

+ 1 - 1
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)
 

+ 12 - 16
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

+ 4 - 4
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],