Update app.py

app.py

@@ -2,6 +2,11 @@ import gradio as gr
 from predict import predict_masks, get_mask_for_label
 import glob
 
+def wrapper_func(input_img_path, selected_label):
+    output_result, output_heading, labels = predict_masks(input_img_path)
+    mask = get_mask_for_label(output_result, selected_label)
+    return mask
+
 demo = gr.Blocks()
 
 with demo:
@@ -28,7 +33,7 @@ with demo:
 
     gr.Markdown("**Examples:**")
     submit_button.click(predict_masks, inputs=[input_image], outputs=[output_mask, output_heading, labels_dropdown])
-    generate_mask_button.click(get_mask_for_label, inputs=[predict_masks(input_image), labels_dropdown], outputs=[selected_mask])
+    generate_mask_button.click(wrapper_func, inputs=[input_image, labels_dropdown], outputs=[selected_mask])
     
     gr.Markdown('\n Demo created by: <a href=\"https://www.linkedin.com/in/theo-belen-halimi/">Théo Belen-Halimi</a>')