diff --git a/detect.py b/detect.py index a8d7e0b6a8c9..eb29022fe80b 100644 --- a/detect.py +++ b/detect.py @@ -8,6 +8,7 @@ import argparse import os +import platform import sys from pathlib import Path @@ -107,7 +108,14 @@ def wrap_frozen_graph(gd, inputs, outputs): elif saved_model: model = tf.keras.models.load_model(w) elif tflite: - interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model + if "edgetpu" in w: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python + import tflite_runtime.interpreter as tflri + delegate = {'Linux': 'libedgetpu.so.1', # install libedgetpu https://coral.ai/software/#edgetpu-runtime + 'Darwin': 'libedgetpu.1.dylib', + 'Windows': 'edgetpu.dll'}[platform.system()] + interpreter = tflri.Interpreter(model_path=w, experimental_delegates=[tflri.load_delegate(delegate)]) + else: + interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model interpreter.allocate_tensors() # allocate input_details = interpreter.get_input_details() # inputs output_details = interpreter.get_output_details() # outputs