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:
    
    gr.Markdown("# **<p align='center'>FurnishAI</p>**")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("**Inputs**")
            input_image = gr.Image(type='filepath',label="Input Image", show_label=True)
            labels_dropdown = gr.Dropdown(label="Labels", show_label=True)
        
        with gr.Column():
            gr.Markdown("**Outputs**")
            output_heading = gr.Textbox(label="Output Type", show_label=True)
            output_mask = gr.Image(label="Predicted Masks", show_label=True)
            selected_mask = gr.Image(label="Selected Mask", show_label=True)
    
    gr.Markdown("**Predict**")
    
    with gr.Row():
        submit_button = gr.Button("Submit")
        generate_mask_button = gr.Button("Generate Mask")

    gr.Markdown("**Examples:**")
    submit_button.click(predict_masks, inputs=[input_image], outputs=[output_mask, output_heading, labels_dropdown])
    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>')

demo.launch(debug=True)