import gradio as gr
from predict import predict_masks
import glob

##Create list of examples to be loaded
example_list = glob.glob("examples/*")
example_list = list(map(lambda el:[el], example_list))

demo = gr.Blocks()

with demo:
    
    gr.Markdown("# **<p align='center'>Mask2Former: Masked Attention Mask Transformer for Universal Segmentation</p>**")
    gr.Markdown("This space demonstrates the use of Mask2Former. It was introduced in the paper [Masked-attention Mask Transformer for Universal Image Segmentation](https://arxiv.org/abs/2112.01527) and first released in [this repository](https://github.com/facebookresearch/Mask2Former/). \
        Before Mask2Former, you'd have to resort to using a specialized architecture designed for solving a particular kind of image segmentation task (i.e. semantic, instance or panoptic segmentation). On the other hand, in the form of Mask2Former, for the first time, we have a single architecture that is capable of solving any segmentation task and performs on par or better than specialized architectures.")
       
    with gr.Box():
        
        
        with gr.Row():
            with gr.Column():
                gr.Markdown("**Inputs**")
                segmentation_task = gr.Dropdown(["semantic", "instance", "panoptic"], value="panoptic", label="Segmentation Task", show_label=True)
                input_image = gr.Image(type='filepath',label="Input Image", 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)
    
    gr.Markdown("**Predict**")
    
    with gr.Box():
        with gr.Row():
            submit_button = gr.Button("Submit")
    
    gr.Markdown("**Examples:**")

    with gr.Column():
        gr.Examples(example_list, [input_image, segmentation_task], [output_mask,output_heading], predict_masks)
        
    
    submit_button.click(predict_masks, inputs=[input_image, segmentation_task], outputs=[output_mask,output_heading])
    
    gr.Markdown('\n Demo created by: <a href=\"https://www.linkedin.com/in/shivalika-singh/\">Shivalika Singh</a>')

demo.launch(debug=True)