Merge pull request #12 from bumble-tech/base64_image

Base64 image input

README.md

@@ -58,6 +58,10 @@ python3 inference.py \
 </code>
 </details>
 
+## Serving
+
+See [Tensorflow Serving example](deployments/tensorflow-serving/README.md)
+
 ## Additional Training
 
 You can finetune the model yourself on your own data, to do so is fairly simple - though you will need the checkpoint files as can be found in `saved_checkpoint/` in `private_detector.zip`

deployments/tensorflow-serving/README.md

@@ -0,0 +1,35 @@
+## Private Detector on Tensorflow Serving
+
+Tensorflow serving image could be pulled via:
+
+```shell
+docker pull vazhega/private-detector:0.2.0
+```
+
+Run: 
+
+```shell
+docker run -p 8501:8501 -e MODEL_NAME=private_detector -t vazhega/private-detector:0.2.0
+```
+
+Ports exposed:
+* REST API: **8501**
+* GRPC: **8502**
+
+Calling a model:
+
+```python
+import base64
+import json
+import requests
+
+image_string = open("./test.jpeg", "rb").read()
+endpoint = "http://localhost:8501/v1/models/private_detector:predict"
+jpeg_bytes = base64.b64encode(image_string).decode('utf-8')
+predict_request = {"instances" : [{"b64": jpeg_bytes}]}
+
+response = requests.post(endpoint, json.dumps(predict_request))
+print(response.json())
+
+>>> {'predictions': [[0.0535223149, 0.946477652]]}
+```

private_detector/private_detector.py

@@ -12,6 +12,7 @@
 from .utils.efficientnet_config import EfficientNetV2Config
 from .utils.effnetv2_model import EffNetV2Model
 from .utils.tensorboard_callback import Callback
+from .utils.preprocess import preprocess_for_evaluation
 
 
 class PrivateDetector():
@@ -558,3 +559,45 @@ def __call__(self, inputs):
             inference_model,
             output_dir
         )
+
+    def save_base64_serving(self, output_dir: str, image_size: int = 480, input_dtype: tf.DType = tf.float16):
+        """
+        Save model as a SavedModel ready for inference, accepting base64 image, returning class probabilities
+
+        Parameters
+        ----------
+        output_dir : str
+            Directory to save model to
+        image_size : int
+            Image size which the image was trained with
+        input_dtype: tf.DType
+            Input dtype image was trained with
+        """
+        image_size = image_size or self.eval_image_size
+
+        # This function will act as a tensorflow Lambda layer
+        def preprocess_input_base64(base64_input_bytes):
+            def decode_bytes(img_bytes):
+                img = tf.image.decode_jpeg(img_bytes, channels=3)
+                img = preprocess_for_evaluation(img, image_size, input_dtype)
+                img = tf.image.resize(img, [image_size, image_size])
+                img = tf.image.convert_image_dtype(img,input_dtype)
+                return img
+
+            base64_input_bytes = tf.reshape(base64_input_bytes, (-1,))
+            return tf.map_fn(lambda img_bytes:
+                             decode_bytes(img_bytes),
+                             elems=base64_input_bytes,
+                             fn_output_signature=input_dtype)
+
+        # Adding base64 input layer and a decoding layers
+        inputs = tf.keras.layers.Input(shape=(), dtype=tf.string, name='b64_input_bytes')
+        decode_inputs = tf.keras.layers.Lambda(preprocess_input_base64, name='decode_image_bytes')(inputs)
+        model_base64 = self.model(decode_inputs)
+
+        # Adding a softmax layer on top to get probabilities as the model's response
+        softmax = tf.keras.layers.Softmax()
+        model_base64_softmax = softmax(model_base64)
+        serving_model = tf.keras.Model(inputs, model_base64_softmax)
+
+        tf.saved_model.save(serving_model, output_dir)

train.py

@@ -214,6 +214,7 @@ def train(train_id: str,
     logger.info(f'Training complete, saving model to {train_id}')
 
     model.save(train_id)
+    model.save_base64_serving(train_id, input_dtype=dtype)
 
     logger.info(f'Model saved to {train_id}')