added application, model, and examples

app.py

@@ -1,12 +1,102 @@
 import gradio as gr
 from huggingface_hub import from_pretrained_keras
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import io
+import base64
+
+
+model = tf.keras.load_model("./tf_model.h5")
 
-model = from_pretrained_keras("Narsil/pet-segmentation")
 
-print(model.summary())
 
 def greet(name):
     return "Hello " + name + "!!"
 
+
+def predict(image):
+    img = np.array(image)
+
+    im = tf.image.resize(img, (128, 128))
+    im = tf.cast(im, tf.float32) / 255.0
+    pred_mask = model.predict(im[tf.newaxis, ...])
+
+    # take the best performing class for each pixel
+    # the output of argmax looks like this [[1, 2, 0], ...]
+    pred_mask_arg = tf.argmax(pred_mask, axis=-1)
+
+    labels = []
+    
+    # convert the prediction mask into binary masks for each class
+    binary_masks = {}
+    mask_codes = {}
+    
+    # when we take tf.argmax() over pred_mask, it becomes a tensor object
+    # the shape becomes TensorShape object, looking like this TensorShape([128]) 
+    # we need to take get shape, convert to list and take the best one
+    
+    rows = pred_mask_arg[0][1].get_shape().as_list()[0]
+    cols = pred_mask_arg[0][2].get_shape().as_list()[0]
+    
+    for cls in range(pred_mask.shape[-1]):
+
+        binary_masks[f"mask_{cls}"] = np.zeros(shape = (pred_mask.shape[1], pred_mask.shape[2])) #create masks for each class
+        
+        for row in range(rows):
+
+            for col in range(cols):
+
+                if pred_mask_arg[0][row][col] == cls:
+                    
+                    binary_masks[f"mask_{cls}"][row][col] = 1
+                else:
+                    binary_masks[f"mask_{cls}"][row][col] = 0
+
+        mask = binary_masks[f"mask_{cls}"]
+        mask *= 255
+        img = Image.fromarray(mask.astype(np.int8), mode="L")
+            
+        # we need to make it readable for the widget
+        with io.BytesIO() as out:
+            img.save(out, format="PNG")
+            png_string = out.getvalue()
+            mask = base64.b64encode(png_string).decode("utf-8")
+
+        mask_codes[f"mask_{cls}"] = mask
+
+        # widget needs the below format, for each class we return label and mask string
+        labels.append({
+            "label": f"LABEL_{cls}",
+            "mask": mask_codes[f"mask_{cls}"],
+            "score": 1.0,
+        })
+
+    return labels["mask"], labels["label"]
+    
+
+
+inputs = gr.inputs.Image(label="Upload a fetal standard plane image", type = 'pil', optional=False)
+outputs = [
+    gr.outputs.Image(label="Segmentation"),
+    gr.outputs.Textbox(type="auto",label="Fetal Plane Prediction")
+]
+
+examples = [
+    "./examples/cat_1.jpg",
+    "./examples/cat_2.jpg",
+    "./examples/dog_1.jpg",
+    "./examples/dog_2.jpg",
+]
+
 iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+iface.launch()
+
+interface = gr.Interface(fn=model.predict, 
+    inputs=inputs,
+    outputs=outputs,
+    # title = title, 
+    # description=description,
+    examples=examples
+    )
+interface.launch()

requirements.txt

@@ -1,2 +1,3 @@
 huggingface_hub
-tensorflow
+tensorflow
+pillow

tf_model.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0258ea75c11d977fae78f747902e48541c5e6996d3d5c700175454ffeb42aa0f
+size 63661584