Update app.py

app.py

@@ -13,45 +13,45 @@ output = gr.outputs.Image()
 def predict(image_input):
   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")
-            return img
+  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")
+      return img