prep

Author: robmarkcole

README.md

@@ -1,5 +1,5 @@
 # tensorflow-lite-rest-server
-Expose tensorflow-lite models via a rest API, and currently object detection is supported. Can be hosted on any of the common platforms including RPi, linux desktop, Mac and Windows.
+Expose tensorflow-lite models via a rest API, and currently object, face & scene detection is supported. Can be hosted on any of the common platforms including RPi, linux desktop, Mac and Windows.
 
 ## Setup
 In this process we create a virtual environment (venv), then install tensorflow-lite [as per these instructions](https://www.tensorflow.org/lite/guide/python) which is platform specific, and finally install the remaining requirements. Note on an RPi (only) it is necessary to manually install pip3, numpy, pillow.
@@ -58,6 +58,13 @@ To detect faces:
 curl -X POST -F image=@tests/faces.jpg 'http://localhost:5000/v1/vision/face'
 ```
 
+To run the scene:
+```
+curl -X POST -F image=@tests/cat.jpg 'http://localhost:5000/v1/vision/scene'
+or
+curl -X POST -F image=@tests/dog.jpg 'http://localhost:5000/v1/vision/scene'
+```
+
 ## Deepstack, Home Assistant & UI
 This API can be used as a drop in replacement for [deepstack object detection](https://github.com/robmarkcole/HASS-Deepstack-object) and [deepstack face detection](https://github.com/robmarkcole/HASS-Deepstack-face) (configuring `detect_only: True`) in Home Assistant. I also created a UI for viewing the predictions of the object detection model [here](https://github.com/robmarkcole/deepstack-ui).
 

helpers.py

@@ -3,9 +3,9 @@
 """
 
 
-def read_coco_labels(file_path):
+def read_labels(file_path):
     """
-    Helper for loading coco_labels.txt
+    Helper for loading labels.txt
     """
     with open(file_path, "r", encoding="utf-8") as f:
         lines = f.readlines()

models/classification/dogs-vs-cats/README.md

@@ -1,5 +1,5 @@
 ## dogs vs cats
 * Custom model trained on [teachable machine image classification](https://teachablemachine.withgoogle.com/train/image)
 * Dataset: 30 cat and 30 dog images from [kaggle dogs vs cats](https://www.kaggle.com/c/dogs-vs-cats)
-* Input size: TBC
+* Input size: 224x224
 * Type: tensorflow lite quantized

tests/cat.jpg

@@ -11,7 +11,7 @@
 import tflite_runtime.interpreter as tflite
 from PIL import Image
 
-from helpers import read_coco_labels
+from helpers import read_labels
 
 app = flask.Flask(__name__)
 
@@ -32,6 +32,10 @@
 OBJ_MODEL = "models/object_detection/mobilenet_ssd_v2_coco/mobilenet_ssd_v2_coco_quant_postprocess.tflite"
 OBJ_LABELS = "models/object_detection/mobilenet_ssd_v2_coco/coco_labels.txt"
 
+SCENE_URL = "/v1/vision/scene"
+SCENE_MODEL = "models/classification/dogs-vs-cats/model.tflite"
+SCENE_LABEL = "models/classification/dogs-vs-cats/labels.txt"
+
 
 @app.route("/")
 def info():
@@ -69,7 +73,9 @@ def predict_face():
         # Process image and get predictions
         face_interpreter.invoke()
         boxes = face_interpreter.get_tensor(face_output_details[0]["index"])[0]
-        classes = face_interpreter.get_tensor(face_output_details[1]["index"])[0]
+        classes = face_interpreter.get_tensor(face_output_details[1]["index"])[
+            0
+        ]
         scores = face_interpreter.get_tensor(face_output_details[2]["index"])[0]
 
         faces = []
@@ -141,7 +147,9 @@ def predict_object():
 
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Flask app exposing tflite models")
+    parser = argparse.ArgumentParser(
+        description="Flask app exposing tflite models"
+    )
     parser.add_argument("--port", default=5000, type=int, help="port number")
     args = parser.parse_args()
 
@@ -152,14 +160,22 @@ def predict_object():
     obj_output_details = obj_interpreter.get_output_details()
     obj_input_height = obj_input_details[0]["shape"][1]
     obj_input_width = obj_input_details[0]["shape"][2]
-    obj_labels = read_coco_labels(OBJ_LABELS)
+    obj_labels = read_labels(OBJ_LABELS)
 
     # Setup face detection
     face_interpreter = tflite.Interpreter(model_path=FACE_MODEL)
     face_interpreter.allocate_tensors()
     face_input_details = face_interpreter.get_input_details()
     face_output_details = face_interpreter.get_output_details()
-    face_input_height = 320
-    face_input_width = 320
+    face_input_height = face_input_details[0]["shape"][1]  # 320
+    face_input_width = face_input_details[0]["shape"][2]  # 320
+
+    # Setup face detection
+    scene_interpreter = tflite.Interpreter(model_path=FACE_MODEL)
+    scene_interpreter.allocate_tensors()
+    scene_input_details = scene_interpreter.get_input_details()
+    scene_output_details = scene_interpreter.get_output_details()
+    scene_input_height = scene_input_details[0]["shape"][1]
+    scene_input_width = scene_input_details[0]["shape"][2]
 
     app.run(host="0.0.0.0", debug=True, port=args.port)