Update app.py

app.py

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