Handle edgetpu model inference (#5372)
Browse files* Handle edgetpu model inference
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Cleanup
Rename `tflite_runtime.interpreter as tflite` to `tflite_runtime.interpreter as tflri` to avoid conflict with existing `tflite` boolean
Co-authored-by: Nam Vu <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <[email protected]>
detect.py
CHANGED
@@ -8,6 +8,7 @@ Usage:
|
|
8 |
|
9 |
import argparse
|
10 |
import os
|
|
|
11 |
import sys
|
12 |
from pathlib import Path
|
13 |
|
@@ -107,7 +108,14 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
|
|
107 |
elif saved_model:
|
108 |
model = tf.keras.models.load_model(w)
|
109 |
elif tflite:
|
110 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
111 |
interpreter.allocate_tensors() # allocate
|
112 |
input_details = interpreter.get_input_details() # inputs
|
113 |
output_details = interpreter.get_output_details() # outputs
|
|
|
8 |
|
9 |
import argparse
|
10 |
import os
|
11 |
+
import platform
|
12 |
import sys
|
13 |
from pathlib import Path
|
14 |
|
|
|
108 |
elif saved_model:
|
109 |
model = tf.keras.models.load_model(w)
|
110 |
elif tflite:
|
111 |
+
if "edgetpu" in w: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
|
112 |
+
import tflite_runtime.interpreter as tflri
|
113 |
+
delegate = {'Linux': 'libedgetpu.so.1', # install libedgetpu https://coral.ai/software/#edgetpu-runtime
|
114 |
+
'Darwin': 'libedgetpu.1.dylib',
|
115 |
+
'Windows': 'edgetpu.dll'}[platform.system()]
|
116 |
+
interpreter = tflri.Interpreter(model_path=w, experimental_delegates=[tflri.load_delegate(delegate)])
|
117 |
+
else:
|
118 |
+
interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model
|
119 |
interpreter.allocate_tensors() # allocate
|
120 |
input_details = interpreter.get_input_details() # inputs
|
121 |
output_details = interpreter.get_output_details() # outputs
|