endpoint handler

handler.py

@@ -0,0 +1,42 @@
+from PIL import Image
+from io import BytesIO
+import base64
+import tensorflow as tf
+
+
+IMAGE_HEIGHT = 28
+IMAGE_WIDTH = 28
+LABELS = ["੦", "੧", "੨", "੩", "੪", "੫", "੬", "੭", "੮", "੯"]
+DEFAULT_TOP_K = 3
+
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.model = tf.keras.models.load_model(path)
+
+    def __call__(self, data):
+        image = Image.open(BytesIO(base64.b64decode(data.pop("inputs").pop("image"))))
+        tensors = self.to_tensors(image)
+        predictions = self.model.predict(tensors)
+        top_k_scores, top_k_label_ids = tf.nn.top_k(predictions, k=data.pop("parameters", {}).pop("topK", DEFAULT_TOP_K))
+
+        return [
+            {
+                "label": LABELS[(int(label_id))],
+                "score": float(score),
+            }
+            for label_id, score in zip(top_k_label_ids[0], top_k_scores[0])
+        ]
+
+    @staticmethod    
+    def to_tensors(image):
+        if image.mode == "RGBA":
+            img = Image.new("RGB", image.size, (255, 255, 255))
+            img.paste(image, mask=image.split()[3])
+            image = img
+        elif image.mode != "RGB":
+            image = image.convert("RGB")
+
+        image = tf.image.resize(image, size=(IMAGE_HEIGHT, IMAGE_WIDTH), antialias=True)
+        image = tf.image.rgb_to_grayscale(image)
+        return tf.reshape(image, shape=(1, IMAGE_HEIGHT, IMAGE_WIDTH))

requirements.txt

@@ -0,0 +1 @@
+tensorflow==2.13.0