tensorflow or tflite exclusively as interpreter

As per bug report ultralytics#5709 I think there should be only one attempt to assign interpreter, and it appears tflite is only ever needed for the case of edgetpu model.
phodgers · Nov 25, 2021 · d16c89b · d16c89b
1 parent f2ca30a
commit d16c89b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/models/common.py b/models/common.py
@@ -337,7 +337,10 @@ def __init__(self, weights='yolov5s.pt', device=None, dnn=True):
             context = model.create_execution_context()
             batch_size = bindings['images'].shape[0]
         else:  # TensorFlow model (TFLite, pb, saved_model)
-            import tensorflow as tf
+            if 'edgetpu' in w.lower():
+                import tflite_runtime.interpreter as tfli
+            else:
+                import tensorflow as tf
             if pb:  # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt
                 def wrap_frozen_graph(gd, inputs, outputs):
                     x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), [])  # wrapped
@@ -354,7 +357,6 @@ def wrap_frozen_graph(gd, inputs, outputs):
             elif tflite:  # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python
                 if 'edgetpu' in w.lower():
                     LOGGER.info(f'Loading {w} for TensorFlow Lite Edge TPU inference...')
-                    import tflite_runtime.interpreter as tfli
                     delegate = {'Linux': 'libedgetpu.so.1',  # install https://coral.ai/software/#edgetpu-runtime
                                 'Darwin': 'libedgetpu.1.dylib',
                                 'Windows': 'edgetpu.dll'}[platform.system()]