Update app.py

app.py

@@ -31,12 +31,6 @@ def process_image(image, prompt, threhsold):
     mask_min = mask.min()
     mask_max = mask.max()
     mask = (mask - mask_min) / (mask_max - mask_min)
-    
-    fig, ax = plt.subplots()
-    ax.imshow(image)
-    ax.imshow(mask, alpha=0.5, cmap="jet")
-    ax.axis("off")
-    plt.tight_layout()
     return mask
 
 
@@ -46,6 +40,12 @@ def get_masks(prompts, img, threhsold):
     for prompt in prompts:
         mask = process_image(img, prompt, threhsold)
         mask = mask > threhsold
+    
+        fig, ax = plt.subplots()
+        ax.imshow(img)
+        ax.imshow(mask, alpha=0.5, cmap="jet")
+        ax.axis("off")
+        plt.tight_layout()
         masks.append(mask)
     return masks