use new pycoral libraries

This commit is contained in:
Blake Blackshear 2021-08-07 15:34:55 -05:00
parent 4818c08fe2
commit dc759a3e56
3 changed files with 17 additions and 21 deletions

View File

@ -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)

View File

@ -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

View File

@ -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],