updated app.py

app.py

@@ -12,8 +12,49 @@ inputs = gr.inputs.Image()
 outputs = gr.outputs.Image()
 
 # inference fn
-def predict(image_input):
-  pass
+def predict(inputs):
+  # 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 = 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")
+      return img
+
   
-gr.Interface(predict, inputs = inputs, outputs = outputs)
+gr.Interface(predict, inputs = inputs, outputs = outputs).launch()