import gradio as gr
from predict import predict_masks, get_mask_for_label
import glob

demo = gr.Blocks()

with demo:
    
    gr.Markdown("# **<p align='center'>FurnishAI</p>**")
    
    with gr.Box():
        
        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.Box():
        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(get_mask_for_label, inputs=[predict_masks(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)