Update pipeline.py

pipeline.py

@@ -13,31 +13,43 @@ from PIL import Image
 
 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
         with Image.open(inputs) as img:
             img = np.array(img)
 
         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]))
-            for row in range(pred_mask_arg[0][1].get_shape().as_list()[0]):
+            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(pred_mask_arg[0][2].get_shape().as_list()[0]):
+                for col in range(cols):
 
                     if pred_mask_arg[0][row][col] == cls:
                         
@@ -48,7 +60,8 @@ class PreTrainedPipeline():
             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()
@@ -56,9 +69,8 @@ class PreTrainedPipeline():
 
             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}"],