Update app.py

app.py

@@ -54,7 +54,7 @@ def perform_prediction(image):
     mask_image = show_mask_on_image(image, predicted_mask[:, max_iou_index], return_image=True)
     return mask_image
 
-#Function to overlay mask on the image
+# Function to overlay mask on the image
 def show_mask_on_image(raw_image, mask, return_image=False):
     if not isinstance(mask, torch.Tensor):
         mask = torch.Tensor(mask)
@@ -62,7 +62,7 @@ def show_mask_on_image(raw_image, mask, return_image=False):
     if len(mask.shape) == 4:
         mask = mask.squeeze()
 
-    fig, axes = plt.subplots(1,1,figsize=(15,15))
+    fig, axes = plt.subplots(1, 1, figsize=(15, 15))
 
     mask = mask.cpu().detach()
     axes.imshow(np.array(raw_image))
@@ -73,9 +73,9 @@ def show_mask_on_image(raw_image, mask, return_image=False):
     if return_image:
         fig = plt.gcf()
         fig.canvas.draw()
-        #Convert plot to image
+        # Convert plot to image
         img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
-        img = img.reshape(fig.canvas.get_width_height()[::-1]+(3,))
+        img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
         img = Image.fromarray(img)
         plt.close(fig)
         return img