fix bugs in tflite backend

adiazulay · adiazulay · commit 58acc30c6044 · 2020-10-22T16:25:12.000-07:00
diff --git a/src/lobe/backends/_backend_tflite.py b/src/lobe/backends/_backend_tflite.py
@@ -16,7 +16,7 @@ class ImageClassificationModel:
     __MAX_UINT8 = 255
 
     def __init__(self, signature):
-        self.__model_path = "{}/{}.tflite".format(
+        self.__model_path = "{}/{}".format(
             signature.model_path, signature.filename
         )
         self.__tflite_predict_fn = None
@@ -54,7 +54,7 @@ def predict(self, image: Image.Image) -> PredictionResult:
         )
 
         confidences_output = self.__tflite_predict_fn.get_tensor(
-            output_details[2]["index"]
+            output_details[1]["index"]
         )
 
         confidences = np.squeeze(confidences_output)

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ class ImageClassificationModel:`
`16`	`16`	`__MAX_UINT8 = 255`
`17`	`17`
`18`	`18`	`def __init__(self, signature):`
`19`		`- self.__model_path = "{}/{}.tflite".format(`
	`19`	`+ self.__model_path = "{}/{}".format(`
`20`	`20`	`signature.model_path, signature.filename`
`21`	`21`	`)`
`22`	`22`	`self.__tflite_predict_fn = None`
`@@ -54,7 +54,7 @@ def predict(self, image: Image.Image) -> PredictionResult:`
`54`	`54`	`)`
`55`	`55`
`56`	`56`	`confidences_output = self.__tflite_predict_fn.get_tensor(`
`57`		`- output_details[2]["index"]`
	`57`	`+ output_details[1]["index"]`
`58`	`58`	`)`
`59`	`59`
`60`	`60`	`confidences = np.squeeze(confidences_output)`