Merge pull request lobe#8 from lobe/'aa-bugfix'

adammenges · web-flow · commit f11605312815 · 2020-10-29T10:32:55.000-07:00
Bug fix and TF Lite  runtime
diff --git a/setup.py b/setup.py
@@ -1,15 +1,18 @@
 from setuptools import setup, find_packages
 setup(
     name="lobe",
-    version="0.1.0",
+    version="0.2.0",
     packages=find_packages("src"),
     package_dir={"": "src"},
     install_requires=[
         "numpy",
         "pillow",
         "requests",
         "tensorflow>=1.15.2,<2;platform_machine!='armv7l'",
-        "tensorflow<1.14.0 ; platform_machine=='armv7l'"
+        "tflite_runtime ; platform_machine=='armv7l'"
     ],
-    dependency_links=["https://www.piwheels.org/simple/tensorflow"]
+    dependency_links=[
+        "https://www.piwheels.org/simple/tensorflow",
+        "https://dl.google.com/coral/python/tflite_runtime-2.1.0.post1-cp37-cp37m-linux_armv7l.whl"
+    ]
 )
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)`