Merge pull request #6 from adiazulay/tflite-patch

adiazulay · web-flow · commit dfec7f956936 · 2020-07-23T15:02:44.000-07:00
Turns on TFLite Support
diff --git a/src/lobe/ImageModel.py b/src/lobe/ImageModel.py
@@ -17,7 +17,7 @@ def load_from_signature(signature: Signature) -> ImageModel:
     model_format = signature.format
     if model_format == "tf":
         from .backends import _backend_tf as backend
-    elif model_format == "tflite":
+    elif model_format == "tf_lite":
         from .backends import _backend_tflite as backend
     else:
         raise ValueError("Model is an unsupported format")
diff --git a/src/lobe/backends/_backend_tflite.py b/src/lobe/backends/_backend_tflite.py
@@ -16,13 +16,12 @@ class ImageClassificationModel:
     __MAX_UINT8 = 255
 
     def __init__(self, signature):
-        self.__model_path = "{}/{}".format(
+        self.__model_path = "{}/{}.tflite".format(
             signature.model_path, signature.filename
         )
         self.__tflite_predict_fn = None
         self.__labels = signature.classes
 
-        raise ImportError("TFLite not yet supported")
 
     def __load(self):
         self.__tflite_predict_fn = tflite.Interpreter(
@@ -55,7 +54,7 @@ def predict(self, image: Image.Image) -> PredictionResult:
         )
 
         confidences_output = self.__tflite_predict_fn.get_tensor(
-            output_details[1]["index"]
+            output_details[2]["index"]
         )
 
         confidences = np.squeeze(confidences_output)

Original file line number	Diff line number	Diff line change
`@@ -16,13 +16,12 @@ class ImageClassificationModel:`
`16`	`16`	`__MAX_UINT8 = 255`
`17`	`17`
`18`	`18`	`def __init__(self, signature):`
`19`		`- self.__model_path = "{}/{}".format(`
	`19`	`+ self.__model_path = "{}/{}.tflite".format(`
`20`	`20`	`signature.model_path, signature.filename`
`21`	`21`	`)`
`22`	`22`	`self.__tflite_predict_fn = None`
`23`	`23`	`self.__labels = signature.classes`
`24`	`24`
`25`		`- raise ImportError("TFLite not yet supported")`
`26`	`25`
`27`	`26`	`def __load(self):`
`28`	`27`	`self.__tflite_predict_fn = tflite.Interpreter(`
`@@ -55,7 +54,7 @@ def predict(self, image: Image.Image) -> PredictionResult:`
`55`	`54`	`)`
`56`	`55`
`57`	`56`	`confidences_output = self.__tflite_predict_fn.get_tensor(`
`58`		`- output_details[1]["index"]`
	`57`	`+ output_details[2]["index"]`
`59`	`58`	`)`
`60`	`59`
`61`	`60`	`confidences = np.squeeze(confidences_output)`