Create new file

app.py

@@ -0,0 +1,71 @@
+from huggingface_hub import from_pretrained_keras
+import gradio as gr 
+import tensorflow as tf 
+import numpy as np 
+
+model = from_pretrained_keras("keras-io/semantic-segmentation")
+
+inputs = gr.inputs.Image()
+output = gr.output.Image()
+
+
+def predict(image_input):
+  pass 
+  
+class PreTrainedPipeline():
+    def __init__(self, path: str):
+        # load the model
+        self.model = keras.models.load_model(os.path.join(path, "tf_model.h5"))
+
+    def __call__(self, inputs: "Image.Image")-> List[Dict[str, Any]]:
+
+        # convert img to numpy array, resize and normalize to make the prediction
+        img = np.array(inputs)
+
+        im = tf.image.resize(img, (128, 128))
+        im = tf.cast(im, tf.float32) / 255.0
+        pred_mask = self.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
+